#!/usr/bin/env python3
"""
scripts/mvp_setup.py — Meridian MVP developer CLI
===================================================

Usage (from Prod/ root):
  python3 scripts/mvp_setup.py runserver          # start FastAPI dev server
  python3 scripts/mvp_setup.py check              # preflight: env, venv, models
  python3 scripts/mvp_setup.py health             # ping /health endpoint
  python3 scripts/mvp_setup.py test               # run test_evaluation.py
  python3 scripts/mvp_setup.py info               # show config summary
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

# ── Paths ──────────────────────────────────────────────────────────────────────
ROOT        = Path(__file__).resolve().parent.parent   # Prod/
BACKEND     = ROOT / "backend"
MODELS_DIR  = BACKEND / "models"
FRONTEND    = ROOT / "frontend"
ENV_FILE    = ROOT / ".env"
TESTS_DIR   = ROOT / "mental-health-ai" / "tests"

# Model artifacts required for diagnosis mode
REQUIRED_MODELS = {
    "MentalRoBERTa ONNX": MODELS_DIR / "mentalroberta_onnx" / "model.onnx",
    "SVM classifier":     MODELS_DIR / "svm_depression.joblib",
    "Feature scaler":     MODELS_DIR / "scaler.joblib",
}

# Venv search order — prefer .venv, fall back to venv
VENV_CANDIDATES = [ROOT / ".venv", ROOT / "venv"]

# ── ANSI colours ───────────────────────────────────────────────────────────────
RESET  = "\033[0m"
BOLD   = "\033[1m"
GREEN  = "\033[92m"
YELLOW = "\033[93m"
RED    = "\033[91m"
CYAN   = "\033[96m"
DIM    = "\033[2m"


def ok(msg: str):    print(f"  {GREEN}✓{RESET}  {msg}")
def warn(msg: str):  print(f"  {YELLOW}⚠{RESET}  {msg}")
def fail(msg: str):  print(f"  {RED}✗{RESET}  {msg}")
def info(msg: str):  print(f"  {CYAN}→{RESET}  {msg}")
def hdr(msg: str):   print(f"\n{BOLD}{msg}{RESET}")
def dim(msg: str):   print(f"  {DIM}{msg}{RESET}")


# ── Helpers ────────────────────────────────────────────────────────────────────

def find_venv() -> Path | None:
    """Return the first venv directory that has a python3 binary."""
    for candidate in VENV_CANDIDATES:
        py = candidate / "bin" / "python3"
        if py.exists():
            return candidate
    return None


def venv_python() -> Path:
    """Return the venv python3, or sys.executable as fallback."""
    venv = find_venv()
    if venv:
        return venv / "bin" / "python3"
    return Path(sys.executable)


def venv_bin(name: str) -> Path:
    """Return a binary from the venv/bin dir."""
    venv = find_venv()
    if venv:
        p = venv / "bin" / name
        if p.exists():
            return p
    return Path(name)   # fallback: rely on PATH


def load_dotenv_simple() -> dict:
    """Minimal .env parser — no third-party deps needed for this script."""
    result = {}
    if not ENV_FILE.exists():
        return result
    for line in ENV_FILE.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, val = line.partition("=")
        result[key.strip()] = val.strip()
    return result


# ── Commands ───────────────────────────────────────────────────────────────────

def cmd_check(args) -> int:
    """
    Preflight check: venv, Python version, .env, API key, model artifacts.
    Returns exit code (0 = all clear, 1 = blocking issues found).
    """
    hdr("Meridian — Preflight Check")
    errors = 0
    env = load_dotenv_simple()

    # ── 1. Venv ─────────────────────────────────────────────────────
    hdr("Virtual environment")
    venv = find_venv()
    if venv:
        ok(f"Found venv at: {venv.relative_to(ROOT)}/")
    else:
        fail("No venv found. Create one:")
        dim("  python3.11 -m venv venv && source venv/bin/activate")
        dim("  pip install -r requirements.txt")
        errors += 1

    # ── 2. Python version ────────────────────────────────────────────
    hdr("Python version")
    py = venv_python()
    result = subprocess.run([str(py), "--version"], capture_output=True, text=True)
    ver = (result.stdout + result.stderr).strip()
    major, minor = sys.version_info.major, sys.version_info.minor
    if major == 3 and minor >= 11:
        ok(f"{ver}  (≥ 3.11 ✓)")
    else:
        warn(f"{ver} — recommend Python 3.11+ for ARM64 wheel compatibility")

    # ── 3. .env file ─────────────────────────────────────────────────
    hdr(".env file")
    if ENV_FILE.exists():
        ok(f".env found ({ENV_FILE.stat().st_size} bytes)")
    else:
        fail(".env not found. Copy the example and fill in your key:")
        dim("  cp .env.example .env")
        errors += 1

    # ── 4. OpenAI API key ────────────────────────────────────────────
    hdr("OpenAI API key")
    api_key = env.get("OPENAI_API_KEY", "") or os.environ.get("OPENAI_API_KEY", "")
    if not api_key:
        fail("OPENAI_API_KEY not set in .env — therapy mode will not work")
        errors += 1
    elif api_key.startswith("sk-") and len(api_key) > 20:
        masked = api_key[:8] + "..." + api_key[-4:]
        ok(f"Key present: {masked}")
    else:
        warn(f"Key looks unusual (doesn't start with 'sk-'): {api_key[:12]}...")

    therapy_model = env.get("THERAPY_MODEL", "gpt-4o-mini")
    info(f"THERAPY_MODEL = {therapy_model}")

    # ── 5. Model artifacts ───────────────────────────────────────────
    hdr("ML model artifacts")
    models_ok = True
    for label, path in REQUIRED_MODELS.items():
        if path.exists():
            size_mb = path.stat().st_size / 1_000_000
            ok(f"{label}: {path.relative_to(ROOT)}  ({size_mb:.1f} MB)")
        else:
            warn(f"{label}: MISSING — {path.relative_to(ROOT)}")
            dim("    Diagnosis mode will degrade gracefully; therapy mode still works.")
            dim("    Run the Colab notebook (notebooks/train_classifiers.ipynb) to generate.")
            models_ok = False

    if models_ok:
        ok("All 3 model artifacts present — diagnosis mode ready")

    # ── 6. Frontend ──────────────────────────────────────────────────
    hdr("Frontend")
    index = FRONTEND / "index.html"
    if index.exists():
        ok(f"index.html found ({index.stat().st_size:,} bytes)")
    else:
        fail(f"frontend/index.html missing")
        errors += 1

    # ── Summary ──────────────────────────────────────────────────────
    hdr("Summary")
    if errors == 0:
        print(f"\n  {GREEN}{BOLD}All checks passed.{RESET}  Run the server with:\n")
        print(f"    python3 scripts/mvp_setup.py runserver\n")
    else:
        print(f"\n  {RED}{BOLD}{errors} blocking issue(s) found.{RESET}  Fix the items above, then rerun:\n")
        print(f"    python3 scripts/mvp_setup.py check\n")

    return 0 if errors == 0 else 1


def cmd_runserver(args) -> int:
    """
    Start the FastAPI development server using the venv's uvicorn.
    Equivalent to: source venv/bin/activate && uvicorn backend.main:app --reload ...
    """
    hdr("Meridian — Starting Dev Server")

    # Quick preflight: warn if key is missing but don't block
    env = load_dotenv_simple()
    api_key = env.get("OPENAI_API_KEY", "") or os.environ.get("OPENAI_API_KEY", "")
    if not api_key:
        warn("OPENAI_API_KEY not set — therapy mode will fail at runtime")
    else:
        ok(f"OPENAI_API_KEY present  |  model={env.get('THERAPY_MODEL', 'gpt-4o-mini')}")

    for label, path in REQUIRED_MODELS.items():
        if path.exists():
            ok(f"{label} found")
        else:
            warn(f"{label} missing — diagnosis mode will show 'not available' message")

    uvicorn = venv_bin("uvicorn")
    host    = getattr(args, "host", "0.0.0.0")
    port    = str(getattr(args, "port", 8000))
    reload  = not getattr(args, "no_reload", False)

    cmd = [str(uvicorn), "backend.main:app", "--host", host, "--port", port]
    if reload:
        cmd.append("--reload")

    print(f"\n  {CYAN}→{RESET}  {' '.join(cmd)}\n")
    print(f"  {DIM}Frontend: http://{host}:{port}{RESET}")
    print(f"  {DIM}Health:   http://{host}:{port}/health{RESET}")
    print(f"  {DIM}Ctrl+C to stop{RESET}\n")

    os.chdir(ROOT)   # uvicorn must run from project root for module imports
    try:
        subprocess.run(cmd, check=True)
    except KeyboardInterrupt:
        print("\n  Server stopped.")
    except subprocess.CalledProcessError as e:
        fail(f"Server exited with code {e.returncode}")
        return e.returncode
    return 0


def cmd_health(args) -> int:
    """Ping the /health endpoint and print the response."""
    import urllib.request, json as _json
    hdr("Meridian — Health Check")
    port = getattr(args, "port", 8000)
    url  = f"http://localhost:{port}/health"
    info(f"GET {url}")
    try:
        with urllib.request.urlopen(url, timeout=5) as resp:
            body = _json.loads(resp.read())
            ok(f"Status {resp.status}  →  {body}")
            return 0
    except Exception as e:
        fail(f"Could not reach server: {e}")
        dim("  Is the server running? Start it with:")
        dim("    python3 scripts/mvp_setup.py runserver")
        return 1


def cmd_test(args) -> int:
    """Run the MentalRoBERTa evaluation test suite."""
    hdr("Meridian — Running Test Suite")
    test_file = TESTS_DIR / "test_evaluation.py"
    if not test_file.exists():
        fail(f"Test file not found: {test_file.relative_to(ROOT)}")
        return 1

    py  = venv_python()
    threshold = getattr(args, "threshold", None)
    cmd = [str(py), str(test_file)]
    if threshold is not None:
        cmd += ["--threshold", str(threshold)]

    info(f"Running: {' '.join(cmd)}")
    print()
    os.chdir(ROOT)
    result = subprocess.run(cmd)
    return result.returncode


def cmd_info(args) -> int:
    """Print current configuration and model status."""
    hdr("Meridian — Configuration")
    env = load_dotenv_simple()

    print(f"\n  {'Key':<22} {'Value'}")
    print(f"  {'─'*22} {'─'*40}")
    for key in ("THERAPY_MODEL", "APP_ENV", "OPENAI_API_KEY"):
        val = env.get(key, os.environ.get(key, "—"))
        if key == "OPENAI_API_KEY" and val and val != "—":
            val = val[:8] + "..." + val[-4:]
        print(f"  {key:<22} {val}")

    hdr("Models")
    for label, path in REQUIRED_MODELS.items():
        status = f"{GREEN}✓ {path.stat().st_size/1e6:.1f} MB{RESET}" if path.exists() else f"{YELLOW}missing{RESET}"
        print(f"  {label:<28} {status}")

    hdr("Paths")
    venv = find_venv()
    info(f"Project root : {ROOT}")
    info(f"Active venv  : {venv or 'not found'}")
    info(f"Python       : {venv_python()}")
    print()
    return 0


# ── Entrypoint ─────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        prog="mvp_setup.py",
        description="Meridian MVP developer CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
commands:
  runserver   Start the FastAPI dev server (uses venv uvicorn)
  check       Preflight: venv, .env, API key, model artifacts
  health      Ping /health endpoint to verify server is running
  test        Run MentalRoBERTa evaluation test suite
  info        Show current config and model status
        """,
    )
    sub = parser.add_subparsers(dest="command")

    # runserver
    p_run = sub.add_parser("runserver", help="Start FastAPI dev server")
    p_run.add_argument("--host",      default="0.0.0.0", help="Bind host (default: 0.0.0.0)")
    p_run.add_argument("--port",      default=8000, type=int, help="Bind port (default: 8000)")
    p_run.add_argument("--no-reload", dest="no_reload", action="store_true",
                       help="Disable auto-reload (use in production)")

    # check
    sub.add_parser("check", help="Preflight checks")

    # health
    p_health = sub.add_parser("health", help="Ping /health endpoint")
    p_health.add_argument("--port", default=8000, type=int)

    # test
    p_test = sub.add_parser("test", help="Run evaluation test suite")
    p_test.add_argument("--threshold", type=float, default=None,
                        help="Classification threshold (default: 0.5)")

    # info
    sub.add_parser("info", help="Show config and model status")

    args = parser.parse_args()

    dispatch = {
        "runserver": cmd_runserver,
        "check":     cmd_check,
        "health":    cmd_health,
        "test":      cmd_test,
        "info":      cmd_info,
    }

    if args.command is None:
        parser.print_help()
        return 0

    fn = dispatch.get(args.command)
    if fn is None:
        parser.print_help()
        return 1

    sys.exit(fn(args))


if __name__ == "__main__":
    main()