"""
BizNode License Validator
=========================
Validates the USB license written by the provisioning tool.

Required files on USB (removable drive):
    registration.key     — SHA256 registration key
    hardware.key         — hardware binding token
    .biznode_license     — signed JSON license

Security layers:
    1. Ed25519 signature verification (server-signed, tamper-proof)
    2. USB hardware fingerprint binding (factory serial — cannot be copied)
    3. Continuous USB presence monitoring (USBWatchdog thread)

Signature: Ed25519 (NaCl/PyNaCl) signed by server private key.

Usage:
    from services.license_validator import validate_license, LicenseResult
    result = validate_license()
    if not result.valid:
        print(result.message)
"""

import os
import sys
import json
import base64
import hashlib
import logging
import subprocess
import threading
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional, Callable

logger = logging.getLogger("biznode.license")

# ---------------------------------------------------------------------------
# Server Ed25519 public key — Base64 encoded.
# Generate once with: python scripts/generate_license_keypair.py
# Then paste the PUBLIC key here and keep the PRIVATE key in your provisioning
# tool (SERVER_PRIVATE_KEY_B64 in provision.py).
# ---------------------------------------------------------------------------
SERVER_PUBLIC_KEY_B64 = os.environ.get("BIZNODE_SERVER_PUBLIC_KEY", "baCTrrAmUjjy/jAJ3th4CTTgjry6Ftfhkh6I/l/5QJM=")

# Required files that must exist on the USB root
REQUIRED_FILES = ["registration.key", "hardware.key", ".biznode_license"]

# Activation API endpoint on your server
ACTIVATION_API_URL = "https://biznode.1bz.biz/api/activate"

# Exit code used when USB is removed at runtime
USB_REMOVED_EXIT_CODE = 77


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------

@dataclass
class LicenseResult:
    valid:            bool
    message:          str
    tier:             Optional[str] = None    # "trial", "basic", "pro"
    order_id:         Optional[str] = None
    expiry:           Optional[str] = None
    usb_path:         Optional[str] = None
    registration_key: Optional[str] = None
    days_remaining:   Optional[int] = None
    dev_mode:         bool = False
    usb_fingerprint:  Optional[str] = None   # SHA-256 of device_serial|volume_serial


# ---------------------------------------------------------------------------
# USB detection
# ---------------------------------------------------------------------------

def find_usb_with_license() -> Optional[str]:
    """
    Scan removable drives for BizNode license files.
    Returns the drive path (e.g. 'F:\\') or None.
    """
    # Windows — use win32api
    try:
        import win32api
        import win32file
        drives = win32api.GetLogicalDriveStrings().split("\000")[:-1]
        for drive in drives:
            if win32file.GetDriveType(drive) == win32file.DRIVE_REMOVABLE:
                if _has_license_files(drive):
                    return drive
    except ImportError:
        pass

    # Linux / Mac fallback — check /media and /mnt
    for base in ["/media", "/mnt", "/Volumes"]:
        if os.path.isdir(base):
            for name in os.listdir(base):
                path = os.path.join(base, name)
                if _has_license_files(path):
                    return path

    # Also check environment override (for testing without USB)
    override = os.environ.get("BIZNODE_LICENSE_PATH", "")
    if override and _has_license_files(override):
        return override

    return None


def _has_license_files(path: str) -> bool:
    """Check if all required license files are present at the given path."""
    return all(os.path.exists(os.path.join(path, f)) for f in REQUIRED_FILES)


# ---------------------------------------------------------------------------
# USB hardware serial — factory-burned, cannot be copied
# ---------------------------------------------------------------------------

def get_usb_device_serial(drive_path: str) -> Optional[str]:
    """
    Get the factory device serial number of the physical USB drive.
    This is burned into the USB controller firmware and cannot be changed
    by formatting, renaming, or using third-party tools.

    Uses PowerShell Get-Partition → Get-Disk pipeline.
    """
    drive_letter = drive_path.rstrip(":\\/")
    if not drive_letter:
        return None

    try:
        r = subprocess.run(
            ["powershell", "-NoProfile", "-Command",
             f"(Get-Partition -DriveLetter '{drive_letter}' | Get-Disk).SerialNumber"],
            capture_output=True, text=True, timeout=10,
        )
        serial = r.stdout.strip()
        if serial and serial != "0" and len(serial) > 2:
            return serial.strip().upper()
    except Exception:
        pass

    # Fallback: WMI query via PowerShell for older Windows
    try:
        ps_cmd = (
            f"Get-CimInstance Win32_DiskDrive | Where-Object {{ $_.InterfaceType -eq 'USB' }} | "
            f"ForEach-Object {{ $disk = $_; "
            f"Get-CimInstance -Query \"ASSOCIATORS OF {{Win32_DiskDrive.DeviceID='$($disk.DeviceID)'}} "
            f"WHERE AssocClass=Win32_DiskDriveToDiskPartition\" | "
            f"ForEach-Object {{ Get-CimInstance -Query \"ASSOCIATORS OF {{Win32_DiskPartition.DeviceID='$($_.DeviceID)'}} "
            f"WHERE AssocClass=Win32_LogicalDiskToPartition\" }} | "
            f"Where-Object {{ $_.DeviceID -eq '{drive_letter}:' }} | "
            f"ForEach-Object {{ $disk.SerialNumber }} }}"
        )
        r = subprocess.run(
            ["powershell", "-NoProfile", "-Command", ps_cmd],
            capture_output=True, text=True, timeout=15,
        )
        serial = r.stdout.strip()
        if serial and serial != "0" and len(serial) > 2:
            return serial.strip().upper()
    except Exception:
        pass

    return None


def get_usb_volume_serial(drive_path: str) -> Optional[str]:
    """
    Get the volume serial number of the USB drive.
    This is set during formatting — less secure than device serial
    but adds a second factor for the fingerprint.
    """
    try:
        import win32api
        info = win32api.GetVolumeInformation(drive_path)
        # info[1] is the volume serial number (int)
        vol_serial = info[1]
        # Convert to unsigned 32-bit hex
        return f"{vol_serial & 0xFFFFFFFF:08X}"
    except Exception:
        pass

    # Fallback for non-Windows
    try:
        drive_letter = drive_path.rstrip(":\\/")
        r = subprocess.run(
            ["powershell", "-NoProfile", "-Command",
             f"(Get-Volume -DriveLetter '{drive_letter}').UniqueId"],
            capture_output=True, text=True, timeout=10,
        )
        uid = r.stdout.strip()
        if uid:
            return uid.upper()
    except Exception:
        pass

    return None


def compute_usb_fingerprint(device_serial: str, volume_serial: str) -> str:
    """
    Compute SHA-256 fingerprint from the USB's device serial + volume serial.
    This fingerprint is embedded in the signed license JSON at provisioning time.
    Copying license files to another USB will produce a different fingerprint.
    """
    raw = f"{device_serial}|{volume_serial}"
    return hashlib.sha256(raw.encode()).hexdigest()


def get_current_usb_fingerprint(drive_path: str) -> Optional[str]:
    """
    Read the current USB hardware fingerprint from the inserted drive.
    Returns the SHA-256 hex string or None if serials couldn't be read.
    """
    device_serial = get_usb_device_serial(drive_path)
    volume_serial = get_usb_volume_serial(drive_path)

    if not device_serial:
        logger.warning(f"Could not read device serial for {drive_path}")
        return None

    if not volume_serial:
        logger.warning(f"Could not read volume serial for {drive_path}")
        # Still compute with empty volume_serial — device serial is the strong factor
        volume_serial = "UNKNOWN"

    return compute_usb_fingerprint(device_serial, volume_serial)


# ---------------------------------------------------------------------------
# Signature verification
# ---------------------------------------------------------------------------

def verify_signature(license_data: dict, signature_b64: str) -> tuple[bool, str]:
    """
    Verify Ed25519 signature using the server public key.
    Returns (ok, reason).
    """
    # Development mode — provisioning tool writes "DEV_SIGNATURE" when no key is set
    if signature_b64 == "DEV_SIGNATURE":
        if not SERVER_PUBLIC_KEY_B64:
            return True, "dev_mode"
        return False, "DEV_SIGNATURE not accepted in production (server public key is set)"

    if not SERVER_PUBLIC_KEY_B64:
        return False, "SERVER_PUBLIC_KEY_B64 not configured in license_validator.py"

    try:
        from nacl.signing import VerifyKey
        from nacl.exceptions import BadSignatureError

        verify_key = VerifyKey(base64.b64decode(SERVER_PUBLIC_KEY_B64))

        # Reconstruct the exact message that was signed (all fields except signature)
        data_to_verify = {k: v for k, v in license_data.items() if k != "signature"}
        message = json.dumps(data_to_verify, sort_keys=True).encode()
        signature = base64.b64decode(signature_b64)

        verify_key.verify(message, signature)
        return True, "ok"

    except ImportError:
        return False, "PyNaCl not installed — run: pip install pynacl"
    except Exception as e:
        return False, f"Signature invalid: {e}"


# ---------------------------------------------------------------------------
# Main validator
# ---------------------------------------------------------------------------

def validate_license() -> LicenseResult:
    """
    Full license validation:
    1. Find USB with license files
    2. Read and parse all 3 files
    3. Check registration_key matches
    4. Check expiry
    5. Verify Ed25519 signature
    6. Verify USB hardware fingerprint (anti-copy protection)
    """

    # Step 1 — find USB
    usb_path = find_usb_with_license()
    if not usb_path:
        return LicenseResult(
            valid=False,
            message=(
                "No BizNode license USB found.\n"
                "Insert your licensed USB drive and restart.\n"
                "Contact biznode@1bz.biz or https://t.me/biznode_bot for support."
            ),
        )

    # Step 2 — read files
    try:
        reg_key = open(
            os.path.join(usb_path, "registration.key"), encoding="utf-8"
        ).read().strip()

        hw_key = open(
            os.path.join(usb_path, "hardware.key"), encoding="utf-8"
        ).read().strip()

        with open(os.path.join(usb_path, ".biznode_license"), encoding="utf-8") as f:
            license_data = json.load(f)

    except Exception as e:
        return LicenseResult(
            valid=False,
            usb_path=usb_path,
            message=f"License files unreadable: {e}",
        )

    # Step 3 — registration key must match
    if license_data.get("registration_key") != reg_key:
        return LicenseResult(
            valid=False,
            usb_path=usb_path,
            message="License mismatch: registration.key does not match .biznode_license.",
        )

    # Step 4 — expiry check
    try:
        expiry_str = license_data["expiry"]
        # Handle both ISO format and string with microseconds
        expiry = datetime.fromisoformat(expiry_str.split(".")[0])
        now    = datetime.utcnow()
        if now > expiry:
            return LicenseResult(
                valid=False,
                usb_path=usb_path,
                order_id=license_data.get("order_id"),
                expiry=expiry_str,
                message=(
                    f"License expired on {expiry.date()}. "
                    "Contact biznode@1bz.biz to renew your license."
                ),
            )
        days_remaining = (expiry - now).days
    except Exception as e:
        return LicenseResult(
            valid=False,
            usb_path=usb_path,
            message=f"License expiry date is invalid: {e}",
        )

    # Step 5 — signature verification
    signature = license_data.get("signature", "")
    sig_ok, sig_reason = verify_signature(license_data, signature)
    if not sig_ok:
        return LicenseResult(
            valid=False,
            usb_path=usb_path,
            message=f"License signature verification failed: {sig_reason}",
        )

    dev_mode = (sig_reason == "dev_mode")

    # Step 6 — USB hardware fingerprint verification (anti-copy)
    expected_fp = license_data.get("usb_hardware_fingerprint")
    current_fp = get_current_usb_fingerprint(usb_path)

    if expected_fp and current_fp:
        if expected_fp != current_fp:
            return LicenseResult(
                valid=False,
                usb_path=usb_path,
                message=(
                    "USB hardware mismatch — license files appear to be copied "
                    "from a different USB device.\n"
                    "Use your original licensed USB drive."
                ),
            )
    elif expected_fp and not current_fp:
        # License has fingerprint but we couldn't read USB serial
        logger.warning("Could not read USB hardware serial for verification. Allowing (degraded mode).")
    # If license has no fingerprint field (legacy license), skip this check

    return LicenseResult(
        valid=True,
        tier="pro",
        message="License valid." + (" [DEV MODE]" if dev_mode else ""),
        order_id=license_data.get("order_id"),
        expiry=expiry_str,
        usb_path=usb_path,
        registration_key=reg_key,
        days_remaining=days_remaining,
        dev_mode=dev_mode,
        usb_fingerprint=current_fp,
    )


# ---------------------------------------------------------------------------
# USB Watchdog — continuous presence monitoring
# ---------------------------------------------------------------------------

class USBWatchdog:
    """
    Background daemon thread that re-validates the USB license every N seconds.
    If the USB is removed or validation fails, calls the on_failure callback
    and terminates the application.

    Usage:
        watchdog = USBWatchdog(interval=30, on_failure=my_shutdown_func)
        watchdog.start()
    """

    def __init__(self, interval: int = 30, on_failure: Optional[Callable] = None):
        self._interval = interval
        self._on_failure = on_failure
        self._stop_event = threading.Event()
        self._thread: Optional[threading.Thread] = None

    def start(self):
        self._thread = threading.Thread(
            target=self._run, daemon=True, name="USBWatchdog"
        )
        self._thread.start()
        logger.info(f"USB watchdog started (checking every {self._interval}s)")

    def stop(self):
        self._stop_event.set()
        if self._thread:
            self._thread.join(timeout=5)

    def _run(self):
        while not self._stop_event.wait(self._interval):
            # Quick check: is the USB still physically present?
            usb_path = find_usb_with_license()
            if not usb_path:
                logger.critical("USB license drive removed! Shutting down...")
                self._trigger_shutdown("USB license drive has been removed.")
                return

            # Full validation (expiry, signature, fingerprint)
            result = validate_license()
            if not result.valid:
                logger.critical(f"USB license check failed: {result.message}")
                self._trigger_shutdown(result.message)
                return

    def _trigger_shutdown(self, reason: str):
        print()
        print("=" * 60)
        print("  BIZNODE — USB LICENSE REMOVED")
        print("=" * 60)
        print()
        print(f"  {reason}")
        print()
        print("  All services will shut down now.")
        print("  Insert your USB and restart with: start.bat")
        print()
        print("=" * 60)
        sys.stdout.flush()

        if self._on_failure:
            try:
                self._on_failure(reason)
            except Exception:
                pass

        # Force exit — os._exit bypasses cleanup but is the only reliable
        # way to kill a process with a running asyncio event loop from a daemon thread
        os._exit(USB_REMOVED_EXIT_CODE)


# ---------------------------------------------------------------------------
# Basic license key validation (BZN-XXXXX-XXXXX-XXXXX)
# ---------------------------------------------------------------------------

BASIC_LICENSE_API = "https://biznode.1bz.biz/api_license.php"
BASIC_KEY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".license_key")


def _save_basic_key(data: dict):
    """Persist Basic license key data to .license_key file."""
    try:
        with open(BASIC_KEY_FILE, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
    except OSError as e:
        logger.warning(f"Could not save license key: {e}")


def _load_basic_key() -> Optional[dict]:
    """Load saved Basic license key from .license_key file."""
    if not os.path.exists(BASIC_KEY_FILE):
        return None
    try:
        with open(BASIC_KEY_FILE, encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def validate_basic_license(key: str = None, fingerprint: str = None) -> LicenseResult:
    """
    Validate a Basic license key (BZN-XXXXX-XXXXX-XXXXX).

    - If key is provided, validates against the server and saves on success.
    - If key is None, loads from .license_key file and checks cached validity.
    - Offline grace: allows 7 days since last server validation.

    Returns LicenseResult with tier="basic" on success.
    """
    import re
    import urllib.request
    import urllib.error

    # Load existing key if none provided
    saved = _load_basic_key()
    if key is None:
        if saved is None:
            return LicenseResult(valid=False, message="No Basic license key found.")
        key = saved.get("license_key", "")

    # Validate format
    if not re.match(r'^BZN-[A-Z0-9]{5}-[A-Z0-9]{5}-[A-Z0-9]{5}$', key):
        return LicenseResult(valid=False, message="Invalid license key format. Expected: BZN-XXXXX-XXXXX-XXXXX")

    # Validate against server
    try:
        body = json.dumps({
            "license_key": key,
            "machine_fingerprint": fingerprint or "",
        }).encode()
        req = urllib.request.Request(
            BASIC_LICENSE_API,
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read())

        if data.get("ok"):
            key_data = {
                "license_key":    key,
                "fingerprint":    fingerprint,
                "validated_at":   datetime.utcnow().isoformat(),
                "expiry":         data.get("expiry", ""),
                "days_remaining": data.get("days_remaining", 365),
                "order_id":       data.get("order_id", ""),
                "tier":           "basic",
            }
            _save_basic_key(key_data)

            return LicenseResult(
                valid=True,
                tier="basic",
                message="Basic license valid.",
                order_id=data.get("order_id"),
                expiry=data.get("expiry"),
                registration_key=key,
                days_remaining=data.get("days_remaining", 365),
            )
        else:
            error = data.get("error", "License validation failed.")
            return LicenseResult(valid=False, message=error)

    except urllib.error.HTTPError as e:
        try:
            err_data = json.loads(e.read())
            return LicenseResult(valid=False, message=err_data.get("error", f"HTTP {e.code}"))
        except Exception:
            return LicenseResult(valid=False, message=f"Server error: HTTP {e.code}")

    except Exception:
        # Network unreachable — check cached validation
        if saved and saved.get("license_key") == key:
            validated_at = saved.get("validated_at", "")
            try:
                last_check = datetime.fromisoformat(validated_at)
                days_since = (datetime.utcnow() - last_check).days
                if days_since <= 7:
                    logger.info(f"Basic license offline grace — last validated {days_since} day(s) ago.")
                    return LicenseResult(
                        valid=True,
                        tier="basic",
                        message=f"Basic license valid (offline, last checked {days_since}d ago).",
                        order_id=saved.get("order_id"),
                        expiry=saved.get("expiry"),
                        registration_key=key,
                        days_remaining=saved.get("days_remaining"),
                    )
                else:
                    return LicenseResult(
                        valid=False,
                        message=f"Basic license needs server re-validation (last checked {days_since} days ago). Connect to internet.",
                    )
            except Exception:
                pass

        return LicenseResult(valid=False, message="Cannot reach license server. Check internet connection.")


# ---------------------------------------------------------------------------
# Online activation helper
# ---------------------------------------------------------------------------

def activate_online(
    registration_key: str,
    hardware_key: str,
    usb_path: str,
    api_url: str = ACTIVATION_API_URL,
) -> tuple[bool, str]:
    """
    Call the activation API to get or renew a signed license.
    Writes the returned license back to USB if successful.
    Returns (ok, message).
    """
    try:
        import requests
        r = requests.post(
            api_url,
            json={"registration_key": registration_key, "hardware_key": hardware_key},
            timeout=15,
        )
        if r.status_code == 404:
            return False, "Registration key not found on server."
        if r.status_code == 403:
            return False, "License revoked or suspended."
        r.raise_for_status()

        data = r.json()
        if not data.get("ok"):
            return False, data.get("error", "Activation failed.")

        license_data = data["license"]

        # Write updated license back to USB
        license_path = os.path.join(usb_path, ".biznode_license")
        with open(license_path, "w", encoding="utf-8") as f:
            json.dump(license_data, f, indent=2)

        return True, "License activated and updated on USB."

    except Exception as e:
        return False, f"Activation failed: {e}"


# ---------------------------------------------------------------------------
# CLI test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("BizNode License Validator\n")
    result = validate_license()
    print(f"Valid       : {result.valid}")
    print(f"Message     : {result.message}")
    print(f"Order ID    : {result.order_id}")
    print(f"Expiry      : {result.expiry}")
    print(f"Days left   : {result.days_remaining}")
    print(f"USB path    : {result.usb_path}")
    print(f"USB finger  : {result.usb_fingerprint}")
    print(f"Dev mode    : {result.dev_mode}")

    # Show current USB hardware info for provisioning
    if result.usb_path:
        print(f"\n--- USB Hardware Info ---")
        ds = get_usb_device_serial(result.usb_path)
        vs = get_usb_volume_serial(result.usb_path)
        print(f"Device serial : {ds}")
        print(f"Volume serial : {vs}")
        if ds:
            fp = compute_usb_fingerprint(ds, vs or "UNKNOWN")
            print(f"Fingerprint   : {fp}")
            print(f"\nAdd this to .biznode_license before signing:")
            print(f'  "usb_hardware_fingerprint": "{fp}"')
