"""
epi dev — zero-friction developer onboarding for EPI.

Design principles:
- Real capture first. A live-recorded case is always preferred over injected demo data.
- Label honestly. "Live captured" vs "[DEMO - simulated]" — never blur the two.
- Force the proof loop. Export + verify is the product. Users must reach it.
- Pipeline visible. Every stage prints so users see EPI working, not appearing.
- No config required. Works with no API key, no internet, no setup.
"""
from __future__ import annotations

import json
import os
import socket
import subprocess
import sys
import time
import uuid
import webbrowser
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.panel import Panel

console = Console()

# ---------------------------------------------------------------------------
# Demo script — uses real EPI SDK path (wrap_client + @record)
# Mock LLM built in so no API key is needed
# ---------------------------------------------------------------------------

_DEMO_SCRIPT = '''\
"""
EPI demo — Refund approval agent.
Generated by `epi dev`.  Edit freely.
"""
import os, time

from epi_recorder import record, wrap_client  # noqa: F401 — wrap_client used in real-API path

# ── Mock LLM (runs with no API key) ──────────────────────────────────────
class _MockMsg:
    role = "assistant"
    content = (
        "Approved. Order ORD-9001 qualifies: Gold customer, $900, "
        "12 days since purchase (within 30-day window). Risk: low."
    )

class _MockChoice:
    finish_reason = "stop"
    message = _MockMsg()

class _MockUsage:
    prompt_tokens = 38
    completion_tokens = 28
    total_tokens = 66

class _MockResponse:
    choices = [_MockChoice()]
    usage = _MockUsage()

class _MockCompletions:
    def create(self, **kwargs):
        time.sleep(0.25)
        return _MockResponse()

class _MockChat:
    completions = _MockCompletions()

class _MockClient:
    chat = _MockChat()


def run_refund_agent():
    with record("demo_refund.epi"):
        from epi_recorder import get_current_session
        session = get_current_session()

        if session:
            session.log_step("agent.run.start", {
                "agent_name": "RefundApprovalAgent",
                "agent_type": "decision",
                "user_input": "Process refund for order ORD-9001 ($900)",
                "goal": "Determine whether to approve or reject the refund",
            })

        if session:
            session.log_step("tool.call", {
                "tool": "lookup_order",
                "input": {"order_id": "ORD-9001"},
            })
            session.log_step("tool.response", {
                "tool": "lookup_order",
                "output": {
                    "order_id": "ORD-9001",
                    "amount_usd": 900,
                    "customer_status": "gold",
                    "days_since_purchase": 12,
                },
                "status": "success",
            })

        # Use real OpenAI if available, otherwise mock
        use_real = False
        try:
            from openai import OpenAI
            if os.getenv("OPENAI_API_KEY"):
                use_real = True
        except ImportError:
            pass

        model = "gpt-4o" if use_real else "mock-llm"
        msgs = [
            {"role": "system", "content": "You are a refund decision agent. Be concise."},
            {"role": "user",   "content": "Order ORD-9001: $900, gold customer, 12 days old. Approve?"},
        ]

        if session:
            session.log_step("llm.request", {"provider": "openai", "model": model, "messages": msgs})

        import time as _time
        _t0 = _time.time()

        if use_real:
            print("[EPI demo] Using real OpenAI API")
            client = wrap_client(OpenAI())
            response = client.chat.completions.create(model=model, messages=msgs)
            answer = response.choices[0].message.content
        else:
            print("[EPI demo] No OPENAI_API_KEY — using built-in mock LLM")
            response = _MockClient().chat.completions.create(model=model, messages=msgs)
            answer = response.choices[0].message.content

        _latency = round(_time.time() - _t0, 3)
        print(f"[EPI demo] Decision: {answer[:80]}")

        if session:
            session.log_step("llm.response", {
                "provider": "openai",
                "model": model,
                "choices": [{"message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}],
                "usage": {
                    "prompt_tokens": getattr(getattr(response, "usage", None), "prompt_tokens", 38),
                    "completion_tokens": getattr(getattr(response, "usage", None), "completion_tokens", 28),
                    "total_tokens": getattr(getattr(response, "usage", None), "total_tokens", 66),
                },
                "latency_seconds": _latency,
            })

        if session:
            session.log_step("agent.decision", {
                "decision": "approve_refund",
                "output": {"order_id": "ORD-9001", "amount_usd": 900, "action": "APPROVE"},
                "confidence": 0.94,
                "rationale": answer,
                "review_required": True,
            })
            session.log_step("agent.run.end", {
                "agent_name": "RefundApprovalAgent",
                "success": True,
            })

        return answer


if __name__ == "__main__":
    run_refund_agent()
'''


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _find_free_port(preferred: int = 8787) -> int:
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("127.0.0.1", preferred))
            return preferred
    except OSError:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("127.0.0.1", 0))
            return s.getsockname()[1]


def _wait_for_port(host: str, port: int, timeout: float = 20.0) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            with socket.create_connection((host, port), timeout=0.5):
                return True
        except OSError:
            time.sleep(0.3)
    return False


def _find_demo_epi(stem: str = "demo_refund") -> Optional[Path]:
    """Find the most recently written .epi file matching stem."""
    candidates: list[Path] = []
    for search_dir in (Path("epi-recordings"), Path(".")):
        candidates.extend(search_dir.glob(f"{stem}*.epi"))
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_mtime)


def _ingest_epi_into_gateway(epi_path: Path, storage_dir: Path) -> Optional[str]:
    """
    Read a local .epi case file and push it into the gateway case store.
    Returns the case_id on success, None on failure.
    Labels the case "[Live captured]" so it's visually distinct from seeded demo data.
    """
    try:
        from epi_gateway.worker import EvidenceWorker

        if not zipfile.is_zipfile(epi_path):
            return None

        with zipfile.ZipFile(epi_path, "r") as zf:
            names = zf.namelist()
            manifest: dict = json.loads(zf.read("manifest.json")) if "manifest.json" in names else {}
            steps: list = []
            if "steps.jsonl" in names:
                raw = zf.read("steps.jsonl").decode("utf-8", errors="replace")
                for line in raw.splitlines():
                    line = line.strip()
                    if line:
                        try:
                            steps.append(json.loads(line))
                        except json.JSONDecodeError:
                            pass

        case_id = (
            manifest.get("recording_id")
            or manifest.get("id")
            or f"live-{uuid.uuid4().hex[:8]}"
        )
        created_at = manifest.get("created_at") or datetime.now(timezone.utc).isoformat()

        worker = EvidenceWorker(storage_dir=storage_dir)
        payload = {
            "id": case_id,
            "title": "[Live captured] Refund ORD-9001 - $900 (Gold customer, 12 days)",
            "created_at": created_at,
            "manifest": manifest,
            "steps": steps,
            "analysis": {
                "review_required": True,
                "summary": (
                    "Agent approved $900 refund for Gold customer. "
                    "Order ORD-9001, 12 days since purchase. Confidence: 94%. "
                    "Awaiting human review."
                ),
            },
        }
        result = worker.upsert_case_payload(payload)
        return result.get("id") or case_id

    except Exception as exc:
        console.print(f"[yellow][!][/yellow] Could not ingest the case file: {exc}")
        return None


def _seed_simulated_case(storage_dir: Path) -> str:
    """
    Inject a clearly-labeled demo case as fallback.
    Always returns a case_id (creates a new UUID if needed).
    """
    case_id = "demo-simulated-refund-9001"
    now_iso = datetime.now(timezone.utc).isoformat()

    try:
        from epi_gateway.worker import EvidenceWorker

        worker = EvidenceWorker(storage_dir=storage_dir)
        payload = {
            "id": case_id,
            "title": "[DEMO - simulated] Refund ORD-9001 - $900 (Gold customer, 12 days)",
            "created_at": now_iso,
            "analysis": {
                "review_required": True,
                "summary": (
                    "Simulated: agent approved $900 refund. "
                    "This case was seeded by `epi dev` because the demo script did not "
                    "produce a live recording. Run `epi dev` again after fixing any errors "
                    "to get a real captured case."
                ),
            },
            "steps": [
                {
                    "kind": "agent.run.start",
                    "content": {
                        "agent_name": "RefundApprovalAgent",
                        "agent_type": "decision",
                        "user_input": "Process refund for order ORD-9001 ($900)",
                        "goal": "Determine whether to approve or reject the refund",
                        "timestamp": now_iso,
                    },
                    "seq": 1,
                    "timestamp": now_iso,
                },
                {
                    "kind": "tool.call",
                    "content": {"tool": "lookup_order", "input": {"order_id": "ORD-9001"}, "timestamp": now_iso},
                    "seq": 2,
                    "timestamp": now_iso,
                },
                {
                    "kind": "tool.response",
                    "content": {
                        "tool": "lookup_order",
                        "output": {"order_id": "ORD-9001", "amount_usd": 900, "customer_status": "gold", "days_since_purchase": 12},
                        "status": "success",
                        "timestamp": now_iso,
                    },
                    "seq": 3,
                    "timestamp": now_iso,
                },
                {
                    "kind": "llm.request",
                    "content": {
                        "provider": "openai",
                        "model": "gpt-4o",
                        "messages": [
                            {"role": "system", "content": "You are a refund decision agent."},
                            {"role": "user",   "content": "Order ORD-9001: $900, gold, 12 days. Approve?"},
                        ],
                        "timestamp": now_iso,
                    },
                    "seq": 4,
                    "timestamp": now_iso,
                },
                {
                    "kind": "llm.response",
                    "content": {
                        "provider": "openai",
                        "model": "gpt-4o",
                        "choices": [{"message": {"role": "assistant", "content": "Approved. Gold customer, within return window. Risk: low."}, "finish_reason": "stop"}],
                        "usage": {"prompt_tokens": 38, "completion_tokens": 28, "total_tokens": 66},
                        "latency_seconds": 0.25,
                        "timestamp": now_iso,
                    },
                    "seq": 5,
                    "timestamp": now_iso,
                },
                {
                    "kind": "agent.decision",
                    "content": {
                        "decision": "approve_refund",
                        "output": {"order_id": "ORD-9001", "amount_usd": 900, "action": "APPROVE"},
                        "confidence": 0.94,
                        "rationale": "Approved. Gold customer, within return window. Risk: low.",
                        "review_required": True,
                        "timestamp": now_iso,
                    },
                    "seq": 6,
                    "timestamp": now_iso,
                },
                {
                    "kind": "agent.run.end",
                    "content": {"agent_name": "RefundApprovalAgent", "success": True, "timestamp": now_iso},
                    "seq": 7,
                    "timestamp": now_iso,
                },
            ],
        }
        worker.upsert_case_payload(payload)
    except Exception as exc:
        console.print(f"[yellow][!][/yellow] Could not seed simulated case: {exc}")

    return case_id


def _run_demo_script(script_path: Path, env: dict) -> bool:
    """Run the demo script. Returns True if it exited with code 0."""
    try:
        result = subprocess.run(
            [sys.executable, str(script_path)],
            env=env,
            timeout=60,
        )
        return result.returncode == 0
    except subprocess.TimeoutExpired:
        console.print("[yellow][!][/yellow] Demo script timed out (60s)")
        return False
    except Exception as exc:
        console.print(f"[yellow][!][/yellow] Demo script error: {exc}")
        return False


def _auto_export_and_verify(storage_dir: Path, case_id: str) -> Optional[Path]:
    """
    Export the case to a .epi case file and run epi verify on it.
    Prints the full verify output so the user sees the signed proof.
    Returns the exported path or None on failure.
    """
    from epi_gateway.worker import EvidenceWorker
    from epi_gateway.main import _build_gateway_signer, _build_settings_from_env

    out_path = Path("refund_case.epi").resolve()

    console.print()
    console.print("[dim]Exporting the case file...[/dim]")

    worker = EvidenceWorker(storage_dir=storage_dir)

    try:
        settings = _build_settings_from_env()
        signer = _build_gateway_signer(settings)
    except Exception:
        signer = None

    try:
        worker.export_case(case_id, out_path, signer_function=signer)
        signed = signer is not None
        console.print(f"[green][OK][/green] Exported: {out_path}")
        if signed:
            console.print("[green][OK][/green] Signed with Ed25519 key")
        else:
            console.print("[yellow][!][/yellow] Unsigned (run `epi keys generate` to enable signing)")
    except Exception as exc:
        console.print(f"[red][FAIL][/red] Export failed: {exc}")
        return None

    console.print()
    console.print("[dim]Running: epi verify refund_case.epi[/dim]")
    console.print()

    result = subprocess.run(
        [sys.executable, "-m", "epi_cli", "verify", str(out_path)],
        capture_output=False,     # let output flow to terminal
        timeout=15,
    )

    return out_path if result.returncode == 0 else None


# ---------------------------------------------------------------------------
# Command
# ---------------------------------------------------------------------------

app = typer.Typer(help="Zero-friction developer onboarding for EPI.")


@app.callback(invoke_without_command=True)
def dev(
    ctx: typer.Context,
    port: int = typer.Option(8787, "--port", "-p", help="Local review workspace port."),
    storage_dir: Path = typer.Option(Path("./evidence_vault"), "--storage-dir", help="Local case storage directory."),
    no_browser: bool = typer.Option(False, "--no-browser", help="Skip opening the browser."),
    demo_script: Path = typer.Option(Path("demo_refund.py"), "--demo-script", help="Path to demo script (auto-generated if missing)."),
    no_run: bool = typer.Option(False, "--no-run", help="Skip running the demo script."),
    no_verify: bool = typer.Option(False, "--no-verify", help="Skip the interactive export + verify step."),
):
    """
    Start a complete EPI demo: refund review, browser case view, export, verify.

    No API key, no internet, no config required.
    """
    if ctx.invoked_subcommand is not None:
        return

    storage_dir = storage_dir.resolve()
    demo_script = demo_script.resolve()

    console.print()
    console.print(Panel.fit(
        "[bold cyan]epi demo[/bold cyan] — refund review in one command\n"
        "[dim]Real capture. Honest labeling. Clear case review.[/dim]",
        border_style="cyan",
    ))
    console.print()

    # ── [0] Generate demo script ──────────────────────────────────────────
    if not demo_script.exists():
        demo_script.write_text(_DEMO_SCRIPT, encoding="utf-8")
        console.print(f"[green]  [+][/green] Generated {demo_script.name}")
    else:
        console.print(f"[dim]  [-] Using existing {demo_script.name}[/dim]")

    # ── [1] Start gateway ─────────────────────────────────────────────────
    actual_port = _find_free_port(port)
    if actual_port != port:
        console.print(f"[yellow]  [!] Port {port} busy — using {actual_port}[/yellow]")

    gateway_url = f"http://127.0.0.1:{actual_port}"

    console.print(f"[dim]  Starting the local review workspace on {gateway_url}...[/dim]")

    gateway_env = os.environ.copy()
    gateway_env["EPI_GATEWAY_STORAGE_DIR"] = str(storage_dir)
    gateway_env["EPI_GATEWAY_RETENTION_MODE"] = "full_content"

    gateway_proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "epi_gateway.main:app",
         "--host", "127.0.0.1", "--port", str(actual_port), "--log-level", "warning"],
        env=gateway_env,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

    if not _wait_for_port("127.0.0.1", actual_port, timeout=20.0):
        console.print("[red]  [FAIL] The local review workspace did not start.[/red]")
        gateway_proc.terminate()
        raise typer.Exit(1)

    console.print(f"[green]  [OK] Local review workspace ready -> {gateway_url}[/green]")

    # ── [2] Run demo script ───────────────────────────────────────────────
    case_id: Optional[str] = None
    case_source = "none"

    if not no_run:
        console.print()
        console.print("[dim]  Running demo_refund.py (capturing decision evidence)...[/dim]")
        console.print()

        run_env = os.environ.copy()
        run_env["EPI_QUIET"] = "0"
        # Write epi file to known location relative to cwd
        run_env["EPI_RECORDINGS_DIR"] = str(Path.cwd() / "epi-recordings")

        script_ok = _run_demo_script(demo_script, env=run_env)

        console.print()
        if script_ok:
            console.print("[green]  [OK] Demo script completed[/green]")
        else:
            console.print("[yellow]  [!] Demo script exited non-zero[/yellow]")

        # ── [3] Ingest the live artifact ──────────────────────────────────
        # Small delay so the .epi file is fully written
        time.sleep(0.5)
        epi_path = _find_demo_epi("demo_refund")

        if epi_path:
            console.print(f"[dim]  Captured: {epi_path.name} ({epi_path.stat().st_size // 1024 + 1} KB)[/dim]")
            console.print("[dim]  Building the review case file...[/dim]")
            case_id = _ingest_epi_into_gateway(epi_path, storage_dir)
            if case_id:
                case_source = "live"
                console.print(f"[green]  [OK] Live captured case ready[/green]")
        else:
            console.print("[yellow]  [!] Could not find demo_refund.epi[/yellow]")

    # ── [4] Fallback: seed simulated case ─────────────────────────────────
    if not case_id:
        console.print("[dim]  Seeding simulated demo case as fallback...[/dim]")
        case_id = _seed_simulated_case(storage_dir)
        case_source = "simulated"
        console.print("[yellow]  [DEMO] Simulated case added (not a real capture)[/yellow]")

    # ── [5] Open browser ──────────────────────────────────────────────────
    console.print()
    if not no_browser:
        time.sleep(0.3)
        webbrowser.open(gateway_url)
        console.print(f"[green]  [OK] Browser opened -> {gateway_url}[/green]")

    # ── [6] Pipeline summary ──────────────────────────────────────────────
    source_label = (
        "[green]Live captured[/green]" if case_source == "live"
        else "[yellow]DEMO - simulated[/yellow]"
    )

    console.print()
    console.rule("[bold cyan]Demo ready[/bold cyan]")
    console.print()
    console.print(f"  [green]✔[/green] Review workspace ready   {gateway_url}")
    if case_source == "live":
        console.print(f"  [green]✔[/green] Decision captured         demo_refund.epi")
    console.print(f"  [green]✔[/green] Case in inbox             Refund ORD-9001 ($900)  ({source_label})")
    console.print(f"  [green]✔[/green] Browser opened            Inbox ready")
    console.print()
    console.print("  [bold]In your browser:[/bold]")
    console.print("    1. Click the case → [bold]Refund ORD-9001[/bold]")
    console.print("    2. Read the agent decision and evidence trail")
    console.print("    3. Click [green]Approve[/green] or [red]Reject[/red]")
    console.print()
    console.rule()

    # ── [7] Interactive export + verify ───────────────────────────────────
    if not no_verify:
        console.print()
        console.print("  [bold]Proof step[/bold] — export this case as a tamper-evident case file.")
        console.print("  This is what you send to auditors, regulators, or downstream systems.")
        console.print()
        try:
            input("  Press Enter to export & verify  (Ctrl+C to skip) → ")
        except (EOFError, KeyboardInterrupt):
            console.print()
            console.print("  [dim]Skipped. Run manually:[/dim]")
            console.print(f"    epi gateway export --case-id {case_id} --out refund_case.epi")
            console.print("    epi verify refund_case.epi")
        else:
            exported = _auto_export_and_verify(storage_dir, case_id)
            if exported:
                console.print()
                console.print(f"  [bold green]✔ Case file verified: {exported.name}[/bold green]")
                console.print("  [dim]This file is cryptographically tamper-evident.[/dim]")
                console.print("  [dim]Share it. Store it. Verify it anywhere.[/dim]")
            console.print()
            console.print("  [dim]To re-verify later:[/dim]")
            console.print("  [dim]  epi verify refund_case.epi[/dim]")

    # ── [8] Persistence note ──────────────────────────────────────────────
    console.print()
    console.rule()
    console.print()
    console.print(f"  Cases stored at:   [cyan]{storage_dir}[/cyan]")
    console.print("  Restart [bold]epi demo[/bold] to continue where you left off.")
    console.print()
    console.print("  Press [bold]Ctrl+C[/bold] to close the local review workspace.")
    console.print()

    # ── [9] Keep gateway alive ────────────────────────────────────────────
    try:
        while True:
            if gateway_proc.poll() is not None:
                console.print("[red][!] The local review workspace exited unexpectedly.[/red]")
                break
            time.sleep(1.0)
    except KeyboardInterrupt:
        pass
    finally:
        console.print()
        console.print("[dim]Stopping the local review workspace...[/dim]")
        try:
            gateway_proc.terminate()
            gateway_proc.wait(timeout=5.0)
        except Exception:
            gateway_proc.kill()
        console.print("[green][OK][/green] Stopped.")
