"""
QuarterBit License Management
=============================

Handles license activation, validation, and usage tracking.

Usage:
    # Activate license (one-time)
    quarterbit activate <LICENSE_KEY>

    # Or in Python
    from quarterbit.license import activate
    activate("QB-XXXX-XXXX-XXXX")
"""

import os
import json
import hashlib
import platform
import urllib.request
import urllib.error
from pathlib import Path
from datetime import datetime, timedelta
from typing import Optional, Dict, Any

# License server endpoint
LICENSE_API = "https://us-central1-quarterbit.cloudfunctions.net"

# Local config directory
CONFIG_DIR = Path.home() / ".quarterbit"
LICENSE_FILE = CONFIG_DIR / "license.json"
USAGE_FILE = CONFIG_DIR / "usage.json"

# License tiers and limits
# Products: axiom (memory-efficient), verity (reproducible)
TIER_LIMITS = {
    "ANONYMOUS": {"gpu_hours_month": 1, "commercial": False, "products": ["axiom"], "price": 0},
    "FREE": {"gpu_hours_month": 5, "commercial": False, "products": ["axiom"], "price": 0},
    "PRO": {"gpu_hours_month": 1000, "commercial": True, "products": ["axiom", "verity"], "price": 49},
    "TEAM": {"gpu_hours_month": 10000, "commercial": True, "products": ["axiom", "verity"], "price": 199},
    "ENTERPRISE": {"gpu_hours_month": float("inf"), "commercial": True, "products": ["axiom", "verity"], "price": "custom"},
}


def _get_machine_id() -> str:
    """Generate a unique machine identifier."""
    data = f"{platform.node()}-{platform.machine()}-{platform.processor()}"
    return hashlib.sha256(data.encode()).hexdigest()[:16]


def _ensure_config_dir():
    """Create config directory if it doesn't exist."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)


def _load_license() -> Optional[Dict[str, Any]]:
    """Load license from local file."""
    if not LICENSE_FILE.exists():
        return None
    try:
        with open(LICENSE_FILE) as f:
            return json.load(f)
    except:
        return None


def _save_license(data: Dict[str, Any]):
    """Save license to local file."""
    _ensure_config_dir()
    with open(LICENSE_FILE, "w") as f:
        json.dump(data, f, indent=2)


def _load_usage() -> Dict[str, Any]:
    """Load usage tracking data."""
    if not USAGE_FILE.exists():
        return {"sessions": [], "gpu_seconds": 0, "month": datetime.now().strftime("%Y-%m")}
    try:
        with open(USAGE_FILE) as f:
            data = json.load(f)
            # Reset if new month
            current_month = datetime.now().strftime("%Y-%m")
            if data.get("month") != current_month:
                data = {"sessions": [], "gpu_seconds": 0, "month": current_month}
            return data
    except:
        return {"sessions": [], "gpu_seconds": 0, "month": datetime.now().strftime("%Y-%m")}


def _save_usage(data: Dict[str, Any]):
    """Save usage tracking data."""
    _ensure_config_dir()
    with open(USAGE_FILE, "w") as f:
        json.dump(data, f, indent=2)


def login() -> Dict[str, Any]:
    """
    Browser-based login flow.
    Opens browser for user to authenticate, then automatically activates.

    Returns:
        dict with activation status and tier info
    """
    import webbrowser
    import time as _time

    machine_id = _get_machine_id()

    # Step 1: Request device auth code
    try:
        req_data = json.dumps({
            "machineId": machine_id,
            "platform": platform.system(),
        }).encode()

        req = urllib.request.Request(
            f"{LICENSE_API}/createDeviceAuth",
            data=req_data,
            headers={"Content-Type": "application/json"},
            method="POST"
        )

        with urllib.request.urlopen(req, timeout=10) as resp:
            result = json.loads(resp.read().decode())

    except urllib.error.URLError:
        raise ValueError("Cannot reach license server. Check your internet connection.")

    device_code = result.get("deviceCode")
    auth_url = result.get("authUrl")

    if not device_code or not auth_url:
        raise ValueError("Failed to initiate login flow")

    # Step 2: Open browser
    print(f"\nOpening browser to complete login...")
    print(f"If browser doesn't open, visit: {auth_url}\n")
    webbrowser.open(auth_url)

    # Step 3: Poll for completion
    print("Waiting for authentication", end="", flush=True)
    for _ in range(60):  # 2 minute timeout
        _time.sleep(2)
        print(".", end="", flush=True)

        try:
            req = urllib.request.Request(
                f"{LICENSE_API}/checkDeviceAuth?code={device_code}",
                method="GET"
            )
            with urllib.request.urlopen(req, timeout=5) as resp:
                status = json.loads(resp.read().decode())

            if status.get("completed"):
                print(" Done!\n")

                # Save license
                license_data = {
                    "key": status.get("licenseKey"),
                    "tier": status.get("tier", "FREE"),
                    "email": status.get("email"),
                    "activated_at": datetime.now().isoformat(),
                    "machine_id": machine_id,
                }
                _save_license(license_data)

                return {
                    "status": "activated",
                    "tier": license_data["tier"],
                    "email": license_data["email"],
                    "message": f"Login successful! Tier: {license_data['tier']}"
                }

        except:
            pass

    print(" Timeout")
    raise ValueError("Login timed out. Please try again.")


def activate(license_key: str) -> Dict[str, Any]:
    """
    Activate a license key.

    Args:
        license_key: License key from quarterbit.dev dashboard

    Returns:
        dict with activation status and tier info

    Raises:
        ValueError: If license key is invalid
    """
    if not license_key or not license_key.startswith("QB-"):
        raise ValueError("Invalid license key format. Expected QB-XXXX-XXXX-XXXX")

    machine_id = _get_machine_id()

    # Validate with server
    try:
        req_data = json.dumps({
            "licenseKey": license_key,
            "machineId": machine_id,
            "platform": platform.system(),
            "pythonVersion": platform.python_version(),
        }).encode()

        req = urllib.request.Request(
            f"{LICENSE_API}/activateLicense",
            data=req_data,
            headers={"Content-Type": "application/json"},
            method="POST"
        )

        with urllib.request.urlopen(req, timeout=10) as resp:
            result = json.loads(resp.read().decode())

    except urllib.error.HTTPError as e:
        error_body = e.read().decode() if e.fp else str(e)
        raise ValueError(f"Activation failed: {error_body}")
    except urllib.error.URLError:
        raise ValueError("Cannot reach license server. Check your internet connection.")

    if not result.get("valid"):
        raise ValueError(result.get("error", "Invalid license key"))

    # Save locally
    license_data = {
        "key": license_key,
        "tier": result.get("tier", "FREE"),
        "email": result.get("email"),
        "activated_at": datetime.now().isoformat(),
        "machine_id": machine_id,
        "expires_at": result.get("expiresAt"),
    }
    _save_license(license_data)

    return {
        "status": "activated",
        "tier": license_data["tier"],
        "email": license_data["email"],
        "message": f"License activated! Tier: {license_data['tier']}"
    }


def deactivate():
    """Remove license from this machine."""
    if LICENSE_FILE.exists():
        LICENSE_FILE.unlink()
    if USAGE_FILE.exists():
        USAGE_FILE.unlink()


def get_license_info() -> Dict[str, Any]:
    """Get current license information."""
    license_data = _load_license()
    usage = _load_usage()
    gpu_hours_used = usage.get("gpu_seconds", 0) / 3600

    if not license_data:
        # No account - anonymous tier (10 hrs)
        limits = TIER_LIMITS["ANONYMOUS"]
        return {
            "tier": "ANONYMOUS",
            "activated": False,
            "limits": limits,
            "usage": {
                "gpu_hours_used": round(gpu_hours_used, 2),
                "gpu_hours_limit": limits["gpu_hours_month"],
                "month": usage.get("month"),
            }
        }

    tier = license_data.get("tier", "FREE")
    limits = TIER_LIMITS.get(tier, TIER_LIMITS["FREE"])

    return {
        "tier": tier,
        "activated": True,
        "email": license_data.get("email"),
        "key": license_data.get("key", "")[:10] + "...",
        "activated_at": license_data.get("activated_at"),
        "limits": limits,
        "usage": {
            "gpu_hours_used": round(gpu_hours_used, 2),
            "gpu_hours_limit": limits["gpu_hours_month"],
            "month": usage.get("month"),
        }
    }


def check_product(product: str) -> bool:
    """
    Check if a specific product is licensed.

    Args:
        product: Product name ("axiom" or "verity")

    Returns:
        True if licensed, False if not
    """
    info = get_license_info()
    tier = info["tier"]
    limits = TIER_LIMITS.get(tier, TIER_LIMITS["FREE"])
    allowed_products = limits.get("products", ["axiom"])
    return product.lower() in allowed_products


def require_product(product: str):
    """
    Require a specific product to be licensed. Raises if not.

    Args:
        product: Product name ("axiom" or "verity")

    Raises:
        LicenseError: If product is not licensed
    """
    if not check_product(product):
        info = get_license_info()
        tier = info["tier"]
        raise LicenseError(
            f"{product.title()} requires PRO tier or higher. "
            f"Current tier: {tier}. Upgrade at quarterbit.dev"
        )


def check_license() -> bool:
    """
    Check if license is valid for current usage.

    Returns:
        True if valid, raises exception if not
    """
    info = get_license_info()
    tier = info["tier"]

    # Free tier always allowed (with limits enforced at usage)
    if tier == "FREE":
        return True

    # Check if license expired (for paid tiers)
    license_data = _load_license()
    if license_data and license_data.get("expires_at"):
        expires = datetime.fromisoformat(license_data["expires_at"].replace("Z", "+00:00"))
        if datetime.now(expires.tzinfo) > expires:
            raise LicenseError("License expired. Please renew at quarterbit.dev")

    return True


def track_usage(gpu_seconds: float):
    """
    Track GPU usage for this session.

    Args:
        gpu_seconds: Number of GPU seconds used
    """
    usage = _load_usage()
    usage["gpu_seconds"] = usage.get("gpu_seconds", 0) + gpu_seconds
    usage["sessions"].append({
        "timestamp": datetime.now().isoformat(),
        "gpu_seconds": gpu_seconds,
    })
    # Keep last 100 sessions only
    usage["sessions"] = usage["sessions"][-100:]
    _save_usage(usage)

    # Check limits
    info = get_license_info()
    gpu_hours_used = usage["gpu_seconds"] / 3600
    gpu_hours_limit = info["limits"]["gpu_hours_month"]

    if gpu_hours_used > gpu_hours_limit:
        tier = info["tier"]
        if tier == "ANONYMOUS":
            # No account - prompt to create free account for more hours
            raise UsageLimitError(
                f"\n{'='*60}\n"
                f"You've used {gpu_hours_limit} GPU-hours this month.\n\n"
                f"Create a FREE account for 5 GPU-hours/month:\n"
                f"  1. Sign up at: quarterbit.dev\n"
                f"  2. Copy your license key from the dashboard\n"
                f"  3. Run: quarterbit activate <YOUR_KEY>\n\n"
                f"Already have an account? Run: quarterbit activate <KEY>\n"
                f"{'='*60}"
            )
        elif tier == "FREE":
            # Has account but hit free limit - prompt to upgrade
            raise UsageLimitError(
                f"\n{'='*60}\n"
                f"You've used {gpu_hours_limit} GPU-hours this month.\n\n"
                f"Upgrade to Pro for 1,000 GPU-hours/month:\n"
                f"  quarterbit.dev/pricing\n\n"
                f"Pro includes: Activation checkpointing, Gradient compression\n"
                f"{'='*60}"
            )


def report_usage():
    """Report usage to server (called periodically)."""
    license_data = _load_license()
    if not license_data:
        return

    usage = _load_usage()

    try:
        req_data = json.dumps({
            "licenseKey": license_data.get("key"),
            "machineId": _get_machine_id(),
            "gpuSeconds": usage.get("gpu_seconds", 0),
            "sessionCount": len(usage.get("sessions", [])),
            "month": usage.get("month"),
        }).encode()

        req = urllib.request.Request(
            f"{LICENSE_API}/reportUsage",
            data=req_data,
            headers={"Content-Type": "application/json"},
            method="POST"
        )

        with urllib.request.urlopen(req, timeout=5) as resp:
            pass  # Fire and forget
    except:
        pass  # Silent fail - usage reporting is best-effort


class LicenseError(Exception):
    """License validation error."""
    pass


class UsageLimitError(Exception):
    """Usage limit exceeded."""
    pass


# CLI interface
def main():
    """CLI entry point for license management."""
    import sys

    if len(sys.argv) < 2:
        print("QuarterBit License Manager")
        print("=" * 40)
        info = get_license_info()
        tier = info['tier']
        print(f"Tier: {tier}")
        print(f"Activated: {info['activated']}")
        if info['activated']:
            print(f"Email: {info.get('email', 'N/A')}")
            print(f"Key: {info.get('key', 'N/A')}")
        print(f"\nUsage this month:")
        print(f"  GPU hours: {info['usage']['gpu_hours_used']:.2f} / {info['usage']['gpu_hours_limit']}")

        if tier == "ANONYMOUS":
            print(f"\nCreate FREE account for 5 GPU-hours/month:")
            print(f"  1. Sign up at quarterbit.dev")
            print(f"  2. Run: quarterbit activate <YOUR_KEY>")

        print(f"\nCommands:")
        print("  quarterbit login           - Login via browser (recommended)")
        print("  quarterbit activate <KEY>  - Activate with license key")
        print("  quarterbit deactivate      - Remove license")
        print("  quarterbit status          - Show this info")
        return

    cmd = sys.argv[1].lower()

    if cmd == "login":
        try:
            result = login()
            print(f"Success! {result['message']}")
            if result.get('email'):
                print(f"Logged in as: {result['email']}")
        except ValueError as e:
            print(f"Error: {e}")
            sys.exit(1)

    elif cmd == "activate":
        if len(sys.argv) < 3:
            print("Usage: quarterbit activate <LICENSE_KEY>")
            sys.exit(1)
        key = sys.argv[2]
        try:
            result = activate(key)
            print(f"Success! {result['message']}")
        except ValueError as e:
            print(f"Error: {e}")
            sys.exit(1)

    elif cmd == "deactivate":
        deactivate()
        print("License deactivated.")

    elif cmd == "status":
        info = get_license_info()
        print(f"Tier: {info['tier']}")
        print(f"Activated: {info['activated']}")
        if info['activated']:
            print(f"Email: {info.get('email', 'N/A')}")
        print(f"GPU hours used: {info['usage']['gpu_hours_used']:.2f} / {info['usage']['gpu_hours_limit']}")

    else:
        print(f"Unknown command: {cmd}")
        sys.exit(1)


if __name__ == "__main__":
    main()
