"""
aquin.inspect / aquin.data_inspect
───────────────────────────────────
Spins up the full Aquin inspection backend locally, opens a tunnel so the
dashboard can reach it, and keeps the registration alive via heartbeats.

Usage
─────
    # Model + data inspection (needs GPU)
    session = aquin.inspect("llama-3.2-1b", api_key="aq-...")
    input("Press Enter to stop…")
    session.stop()

    # Data-only inspection (CPU is fine)
    session = aquin.data_inspect(api_key="aq-...")
    input("Press Enter to stop…")
    session.stop()
"""

from __future__ import annotations

import atexit
import os
import shutil
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Optional

import requests

# ── Defaults ──────────────────────────────────────────────────────────────────

_DEFAULT_BASE_URL = "https://api.aquin.app"
_DEFAULT_PORT     = 7863
_HEARTBEAT_EVERY  = 30   # seconds
_SAE_CACHE_DIR    = Path.home() / ".aquin" / "sae"

# ── Tunnel helpers ────────────────────────────────────────────────────────────

def _start_cloudflared_tunnel(port: int) -> str:
    """Start cloudflared quick-tunnel and return the public URL."""
    proc = subprocess.Popen(
        ["cloudflared", "tunnel", "--url", f"http://localhost:{port}"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    deadline = time.time() + 30
    while time.time() < deadline:
        line = proc.stdout.readline()
        if not line:
            time.sleep(0.1)
            continue
        # cloudflared prints the URL as: https://xxx.trycloudflare.com
        if "trycloudflare.com" in line:
            url = line.strip().split()[-1]
            if url.startswith("http"):
                return url, proc
    proc.terminate()
    raise RuntimeError("[aquin inspect] cloudflared tunnel failed to start within 30s")


def _start_localtunnel(port: int) -> tuple[str, subprocess.Popen]:
    """Start localtunnel (npx lt) and return the public URL."""
    proc = subprocess.Popen(
        ["npx", "--yes", "localtunnel", "--port", str(port)],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    deadline = time.time() + 60
    while time.time() < deadline:
        line = proc.stdout.readline()
        if not line:
            time.sleep(0.1)
            continue
        if "your url is:" in line.lower():
            url = line.strip().split()[-1]
            return url, proc
        if "https://" in line:
            # some versions just print the URL directly
            for part in line.split():
                if part.startswith("https://"):
                    return part, proc
    proc.terminate()
    raise RuntimeError("[aquin inspect] localtunnel failed to start within 60s")


def _open_tunnel(port: int) -> tuple[str, subprocess.Popen]:
    if shutil.which("cloudflared"):
        return _start_cloudflared_tunnel(port)
    try:
        return _start_localtunnel(port)
    except Exception as e:
        raise RuntimeError(
            "[aquin inspect] No tunnel binary found. "
            "Install cloudflared (https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/downloads/) "
            "or Node.js (for localtunnel via npx)."
        ) from e


# ── SAE weight pull ───────────────────────────────────────────────────────────

def _pull_sae_weights(model_id: str, api_key: str, base_url: str) -> Optional[Path]:
    """Download SAE weights for model_id to ~/.aquin/sae/ if not already cached."""
    _SAE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    safe = model_id.replace("/", "_").replace(":", "_")
    dest = _SAE_CACHE_DIR / f"sae_{safe}.pt"

    # Also check for .safetensors variant
    dest_safe = _SAE_CACHE_DIR / f"sae_{safe}.safetensors"
    if dest.exists():
        return dest
    if dest_safe.exists():
        return dest_safe

    print(f"[aquin inspect] pulling SAE weights for {model_id} …", flush=True)
    try:
        resp = requests.get(
            f"{base_url}/api/worker/sae-weights",
            params={"model_id": model_id},
            headers={"Authorization": f"Bearer {api_key}"},
            stream=True,
            timeout=300,
        )
        if resp.status_code == 404:
            print(f"[aquin inspect] SAE weights not available for {model_id} on server — will try local path.", flush=True)
            return None
        resp.raise_for_status()
        with open(dest, "wb") as f:
            for chunk in resp.iter_content(chunk_size=1024 * 1024):
                f.write(chunk)
        print(f"[aquin inspect] SAE weights cached at {dest}", flush=True)
        return dest
    except Exception as e:
        print(f"[aquin inspect] WARNING: could not pull SAE weights: {e}", flush=True)
        return None


# ── Server launcher ───────────────────────────────────────────────────────────

def _find_server_py(server_dir: Optional[Path] = None) -> Path:
    """Locate server.py — bundled inside the package at aquin/server/server.py."""
    candidates = []
    # 1. Bundled inside the installed wheel (aquin/server/server.py)
    candidates.append(Path(__file__).parent / "server" / "server.py")
    # 2. Explicit override
    if server_dir:
        candidates.append(Path(server_dir) / "server.py")
        candidates.append(Path(server_dir) / "inspection-backend" / "server.py")
    # 3. Monorepo dev layout
    candidates.append(Path(__file__).parent.parent.parent / "inspection-backend" / "server.py")
    for p in candidates:
        if p.exists():
            return p
    raise FileNotFoundError(
        "[aquin inspect] Could not find server.py — the package may be corrupted. "
        "Try: pip install --force-reinstall 'aquin[inspect]'"
    )


def _launch_server(port: int, model_id: Optional[str], data_only: bool,
                   api_key: str, base_url: str, openai_api_key: Optional[str]) -> subprocess.Popen:
    server_py = _find_server_py()
    env = {
        **os.environ,
        "LOCAL_MODE":     "1",
        "DATA_ONLY":      "1" if data_only else "0",
        "LOCAL_MODEL_ID": model_id or "llama-3.2-1b",
        # Supabase env vars must already be set in the environment or .env alongside server.py
    }
    if openai_api_key:
        env["OPENAI_API_KEY"] = openai_api_key

    cmd = [
        sys.executable, "-m", "uvicorn",
        "server:app",
        "--host", "0.0.0.0",
        "--port", str(port),
        "--log-level", "warning",
    ]
    proc = subprocess.Popen(
        cmd,
        cwd=str(server_py.parent),
        env=env,
    )

    # Wait for server to be ready
    deadline = time.time() + 120
    while time.time() < deadline:
        try:
            r = requests.get(f"http://localhost:{port}/health", timeout=2)
            if r.status_code == 200:
                info = r.json()
                print(f"[aquin inspect] server ready — model={info.get('model_id')} sae={info.get('sae_loaded')}", flush=True)
                return proc
        except Exception:
            pass
        if proc.poll() is not None:
            raise RuntimeError("[aquin inspect] server process exited unexpectedly during startup")
        time.sleep(1)

    proc.terminate()
    raise RuntimeError("[aquin inspect] server did not become ready within 120s")


# ── Registration + heartbeat ──────────────────────────────────────────────────

def _register(tunnel_url: str, model_id: Optional[str], capabilities: list[str],
               api_key: str, base_url: str):
    requests.post(
        f"{base_url}/api/worker/register-local-server",
        json={"tunnel_url": tunnel_url, "model_id": model_id, "capabilities": capabilities},
        headers={"Authorization": f"Bearer {api_key}"},
        timeout=10,
    ).raise_for_status()


def _deregister(api_key: str, base_url: str):
    try:
        requests.delete(
            f"{base_url}/api/worker/register-local-server",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=5,
        )
    except Exception:
        pass


def _heartbeat_loop(api_key: str, base_url: str, stop_event: threading.Event):
    while not stop_event.wait(_HEARTBEAT_EVERY):
        try:
            requests.patch(
                f"{base_url}/api/worker/register-local-server",
                headers={"Authorization": f"Bearer {api_key}"},
                timeout=5,
            )
        except Exception:
            pass


# ── LocalSession ──────────────────────────────────────────────────────────────

class LocalSession:
    """
    Handle returned by aquin.inspect() / aquin.data_inspect().
    Call .stop() to shut down the server, tunnel, and deregister.
    """

    def __init__(
        self,
        server_proc:   subprocess.Popen,
        tunnel_proc:   subprocess.Popen,
        tunnel_url:    str,
        model_id:      Optional[str],
        api_key:       str,
        base_url:      str,
        stop_event:    threading.Event,
        hb_thread:     threading.Thread,
    ):
        self._server   = server_proc
        self._tunnel   = tunnel_proc
        self.tunnel_url = tunnel_url
        self.model_id   = model_id
        self._api_key   = api_key
        self._base_url  = base_url
        self._stop      = stop_event
        self._hb        = hb_thread
        self._stopped   = False

        atexit.register(self.stop)
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                signal.signal(sig, lambda *_: self.stop())
            except (OSError, ValueError):
                pass

    def stop(self):
        if self._stopped:
            return
        self._stopped = True
        print("[aquin inspect] shutting down …", flush=True)
        self._stop.set()
        self._hb.join(timeout=5)
        _deregister(self._api_key, self._base_url)
        try:
            self._tunnel.terminate()
        except Exception:
            pass
        try:
            self._server.terminate()
            self._server.wait(timeout=10)
        except Exception:
            pass
        print("[aquin inspect] stopped.", flush=True)

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.stop()


# ── Public API ────────────────────────────────────────────────────────────────

def inspect(
    model_id: str = "llama-3.2-1b",
    api_key: str = "",
    base_url: str = _DEFAULT_BASE_URL,
    port: int = _DEFAULT_PORT,
    sae_cache_dir: Optional[Path] = None,
    openai_api_key: Optional[str] = None,
) -> LocalSession:
    """
    Start the full Aquin inspection server locally (model + data) and register
    it with the dashboard.

    Parameters
    ----------
    model_id:       Aquin model key, e.g. "llama-3.2-1b" or "gpt2-small-np"
    api_key:        Your Aquin API key (aq-...)
    base_url:       Aquin API base URL (default: https://api.aquin.app)
    port:           Local port for the inspection server (default: 7863)
    sae_cache_dir:  Where to store downloaded SAE weights (default: ~/.aquin/sae/)
    openai_api_key: Optional OpenAI key for feature labeling
    """
    if not api_key.startswith("aq-"):
        raise ValueError("api_key must start with 'aq-'")

    global _SAE_CACHE_DIR
    if sae_cache_dir:
        _SAE_CACHE_DIR = Path(sae_cache_dir)

    # Pull SAE weights in background — server startup is the bottleneck anyway
    sae_thread = threading.Thread(
        target=_pull_sae_weights,
        args=(model_id, api_key, base_url),
        daemon=True,
    )
    sae_thread.start()

    print(f"[aquin inspect] starting inspection server on port {port} …", flush=True)
    server_proc = _launch_server(port, model_id, data_only=False,
                                 api_key=api_key, base_url=base_url,
                                 openai_api_key=openai_api_key)

    # Wait for SAE pull to finish before opening tunnel so the server is fully ready
    sae_thread.join()

    print("[aquin inspect] opening tunnel …", flush=True)
    tunnel_url, tunnel_proc = _open_tunnel(port)
    print(f"[aquin inspect] tunnel live at {tunnel_url}", flush=True)

    capabilities = ["model", "data"]
    _register(tunnel_url, model_id, capabilities, api_key, base_url)
    print(f"[aquin inspect] registered — open the Aquin dashboard and click Connect.", flush=True)

    stop_event = threading.Event()
    hb_thread = threading.Thread(
        target=_heartbeat_loop,
        args=(api_key, base_url, stop_event),
        daemon=True,
    )
    hb_thread.start()

    return LocalSession(
        server_proc=server_proc,
        tunnel_proc=tunnel_proc,
        tunnel_url=tunnel_url,
        model_id=model_id,
        api_key=api_key,
        base_url=base_url,
        stop_event=stop_event,
        hb_thread=hb_thread,
    )


def data_inspect(
    api_key: str = "",
    base_url: str = _DEFAULT_BASE_URL,
    port: int = _DEFAULT_PORT,
) -> LocalSession:
    """
    Start the Aquin inspection server in data-only mode (no model/SAE loading).
    Suitable for CPU-only machines. All dataset analysis endpoints are available:
    scan, PII, toxicity, bias, synthetic, copyright, opt-out, poisoning, backdoor
    triggers, liability chain, compliance report, audit trail.

    Parameters
    ----------
    api_key:   Your Aquin API key (aq-...)
    base_url:  Aquin API base URL (default: https://api.aquin.app)
    port:      Local port for the inspection server (default: 7863)
    """
    if not api_key.startswith("aq-"):
        raise ValueError("api_key must start with 'aq-'")

    print(f"[aquin inspect] starting data-only inspection server on port {port} …", flush=True)
    server_proc = _launch_server(port, model_id=None, data_only=True,
                                 api_key=api_key, base_url=base_url,
                                 openai_api_key=None)

    print("[aquin inspect] opening tunnel …", flush=True)
    tunnel_url, tunnel_proc = _open_tunnel(port)
    print(f"[aquin inspect] tunnel live at {tunnel_url}", flush=True)

    capabilities = ["data"]
    _register(tunnel_url, None, capabilities, api_key, base_url)
    print("[aquin inspect] registered — open the Aquin dashboard and click Connect.", flush=True)

    stop_event = threading.Event()
    hb_thread = threading.Thread(
        target=_heartbeat_loop,
        args=(api_key, base_url, stop_event),
        daemon=True,
    )
    hb_thread.start()

    return LocalSession(
        server_proc=server_proc,
        tunnel_proc=tunnel_proc,
        tunnel_url=tunnel_url,
        model_id=None,
        api_key=api_key,
        base_url=base_url,
        stop_event=stop_event,
        hb_thread=hb_thread,
    )
