#!/usr/bin/env python3
from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Any, Dict, List

CURRENT_DIR = Path(__file__).resolve().parent
SCRIPTS_DIR = CURRENT_DIR.parent
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

from lab import attack, baseline, browser, catalog, evidence, provision, render, repro, seed, task_queue, validators  # noqa: E402
from lab.config import ADVISORIES_DIR, CASE_RUNS_DIR, ENV_PROFILES_DIR, RUNS_DIR  # noqa: E402
from lab.utils import command_available, ensure_dir, isoformat, load_json_dir, now_utc, read_json, read_yaml, write_json  # noqa: E402


def _load_advisory(canonical_id: str) -> Dict[str, Any]:
    advisory = read_json(ADVISORIES_DIR / f"{canonical_id}.json", default=None)
    if not advisory:
        raise ValueError(f"Unknown advisory: {canonical_id}")
    return advisory


def _run_dir(run_id: str) -> Path:
    path = CASE_RUNS_DIR / run_id
    ensure_dir(path)
    ensure_dir(path / "logs")
    ensure_dir(path / "assets")
    return path


def _compose_run_id(advisory: Dict[str, Any]) -> str:
    return f"{advisory['system_id']}-{advisory['canonical_id']}-{now_utc().strftime('%Y%m%d%H%M%S')}"


def _resolve_profile(advisory: Dict[str, Any]) -> Dict[str, Any]:
    profile = repro.resolve_profile(advisory["canonical_id"], advisory)
    current_profile = read_yaml(ENV_PROFILES_DIR / "core" / advisory["system_id"] / "current.yaml", default={}) or {}
    merged = dict(current_profile)
    merged.update(profile)
    if current_profile.get("services") and not merged.get("services"):
        merged["services"] = current_profile["services"]
    if current_profile.get("baseline_urls") and not merged.get("baseline_urls"):
        merged["baseline_urls"] = current_profile["baseline_urls"]
    if current_profile.get("artifact_mode") and not merged.get("artifact_mode"):
        merged["artifact_mode"] = current_profile["artifact_mode"]
    if current_profile.get("verification_mode") and not merged.get("verification_mode"):
        merged["verification_mode"] = current_profile["verification_mode"]
    if current_profile.get("browser_required"):
        merged.setdefault("browser_assertions", {})
        merged["browser_assertions"].setdefault("required", current_profile["browser_required"])
    if not profile.get("system_id"):
        merged["system_id"] = advisory["system_id"]
    if not profile.get("profile_id"):
        merged["profile_id"] = advisory["canonical_id"]
    return merged


def _build_run_bundle(
    advisory: Dict[str, Any],
    profile: Dict[str, Any],
    run_id: str,
    verification_status: str,
    verification_mode: str,
    artifact_mode: str,
    baseline_refs: List[str],
    attack_steps: List[Dict[str, Any]],
    browser_refs: List[str],
    container_log_refs: List[str],
    request_log_refs: List[str],
    blocked_reason: str | None,
) -> Dict[str, Any]:
    return {
        "run_id": run_id,
        "system_id": advisory["system_id"],
        "advisory_id": advisory["canonical_id"],
        "repro_profile_id": profile["profile_id"],
        "verification_status": verification_status,
        "verification_mode": verification_mode,
        "artifact_mode": artifact_mode,
        "target_env": "local-docker",
        "compose_services": sorted(profile.get("services", {}).keys()),
        "baseline_refs": baseline_refs,
        "attack_steps": attack_steps,
        "browser_refs": browser_refs,
        "container_log_refs": container_log_refs,
        "request_log_refs": request_log_refs,
        "timeline": [],
        "started_at": isoformat(now_utc()),
        "finished_at": isoformat(now_utc()),
        "blocked_reason": blocked_reason,
    }


def cmd_catalog_sync(args) -> int:
    summary = catalog.sync_catalog(write_profiles=True, write_repro_map=True)
    print(summary)
    return 0


def cmd_compose_generate(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    run_dir = _run_dir(args.run_id or f"compose-{advisory['canonical_id']}")
    compose_result = provision.prepare(profile, run_dir, dry_run=True)
    print(compose_result)
    return 0


def cmd_provision(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    run_dir = _run_dir(args.run_id or _compose_run_id(advisory))
    result = provision.prepare(profile, run_dir, dry_run=args.dry_run)
    print(result)
    return 0


def cmd_seed(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    print({"steps": seed.run_seed(profile)})
    return 0


def cmd_baseline(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    run_dir = _run_dir(args.run_id or _compose_run_id(advisory))
    result = baseline.collect(profile, run_dir)
    print(result)
    return 0


def cmd_attack(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    run_dir = _run_dir(args.run_id or _compose_run_id(advisory))
    result = attack.run_attack(profile, advisory, run_dir, dry_run=args.dry_run)
    print(result)
    return 0


def cmd_verify(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    browser_required = bool(profile.get("browser_assertions", {}).get("required"))
    payload = {
        "advisory": advisory["canonical_id"],
        "profile_id": profile["profile_id"],
        "browser_required": browser_required,
        "result": "ready-for-run",
    }
    print(payload)
    return 0


def cmd_run_case(args) -> int:
    advisory = _load_advisory(args.case)
    profile = _resolve_profile(advisory)
    run_id = args.run_id or _compose_run_id(advisory)
    run_dir = _run_dir(run_id)

    provision_result = provision.prepare(profile, run_dir, dry_run=args.dry_run)
    allow_runtime_steps = provision_result.get("status") not in {"blocked-artifact"}
    baseline_payload = (
        baseline.collect(profile, run_dir) if profile.get("baseline_urls") and allow_runtime_steps else {"observations": []}
    )
    attack_payload = (
        attack.run_attack(profile, advisory, run_dir, dry_run=args.dry_run) if allow_runtime_steps else {"steps": []}
    )

    browser_payload = {"required": bool(profile.get("browser_assertions", {}).get("required")), "present": False, "refs": []}
    blocked_reason = provision_result.get("blocked_reason")
    if browser_payload["required"] and not args.dry_run and profile.get("baseline_urls") and allow_runtime_steps:
        browser_payload = browser.capture(profile["baseline_urls"][0], run_dir, prefix="proof")
        if not browser_payload.get("present"):
            blocked_reason = blocked_reason or browser_payload.get("reason")

    compose_path = Path(provision_result["compose_path"])
    container_logs = evidence.collect_container_logs(run_dir, compose_path) if compose_path.exists() and allow_runtime_steps else []

    verification_status = "triage-manual"
    verification_mode = profile.get("verification_mode", "synthetic")
    artifact_mode = profile.get("artifact_mode", profile.get("provisioning_mode", "synthetic"))
    if args.dry_run:
        verification_status = "triage-manual"
        blocked_reason = blocked_reason or "dry-run only"
    elif provision_result.get("status") == "blocked-artifact":
        verification_status = "blocked-artifact"
    elif browser_payload.get("required") and not browser_payload.get("present"):
        verification_status = "triage-manual"
    elif artifact_mode == "synthetic":
        verification_status = "verified-synthetic"
    else:
        verification_status = "verified-real"

    bundle = _build_run_bundle(
        advisory=advisory,
        profile=profile,
        run_id=run_id,
        verification_status=verification_status,
        verification_mode=verification_mode,
        artifact_mode=artifact_mode,
        baseline_refs=[str(run_dir / "logs" / "baseline.json")] if baseline_payload.get("observations") else [],
        attack_steps=attack_payload.get("steps", []),
        browser_refs=browser_payload.get("refs", []),
        container_log_refs=container_logs,
        request_log_refs=[str(run_dir / "logs" / "attack.json"), str(run_dir / "logs" / "baseline.json")],
        blocked_reason=blocked_reason,
    )
    report_refs = render.render_run(bundle)
    bundle["report_refs"] = report_refs
    evidence.write_run_bundle(run_dir, bundle)
    ensure_dir(RUNS_DIR)
    write_json(RUNS_DIR / f"{run_id}.json", bundle)
    render.render_dashboard()
    print(bundle)
    return 0


def cmd_run_system(args) -> int:
    advisories = [item for item in load_json_dir(ADVISORIES_DIR) if item.get("system_id") == args.system]
    selected = advisories[: args.limit]
    for advisory in selected:
        cmd_run_case(argparse.Namespace(case=advisory["canonical_id"], run_id=None, dry_run=args.dry_run))
    print({"system": args.system, "count": len(selected)})
    return 0


def cmd_run_batch(args) -> int:
    if args.from_queue:
        items = task_queue.dequeue(limit=args.limit)
    else:
        task_queue.enqueue_from_registry(only_hotlane=args.only_hotlane, limit=args.limit)
        items = task_queue.dequeue(limit=args.limit)
    for item in items:
        cmd_run_case(argparse.Namespace(case=item["advisory_id"], run_id=None, dry_run=args.dry_run))
    print({"processed": len(items)})
    return 0


def cmd_render_run(args) -> int:
    run = read_json(RUNS_DIR / f"{args.run_id}.json", default=None)
    if not run:
        raise ValueError(f"Unknown run: {args.run_id}")
    print(render.render_run(run))
    return 0


def cmd_serve_dashboard(args) -> int:
    render.render_dashboard()
    import http.server
    import socketserver

    os_dir = str(render.DASHBOARD_DIR if hasattr(render, "DASHBOARD_DIR") else "")
    if not os_dir:
        from lab.config import DASHBOARD_DIR

        os_dir = str(DASHBOARD_DIR)
    handler = http.server.SimpleHTTPRequestHandler
    with socketserver.TCPServer(("127.0.0.1", args.port), handler) as httpd:
        print(f"serving dashboard at http://127.0.0.1:{args.port}/")
        import os

        os.chdir(os_dir)
        httpd.serve_forever()


def cmd_cleanup(args) -> int:
    run = read_json(RUNS_DIR / f"{args.run_id}.json", default=None)
    if not run:
        raise ValueError(f"Unknown run: {args.run_id}")
    compose_path = Path(run["report_refs"]["bundle_dir"]) / "compose" / "compose.yaml"
    if command_available("docker") and compose_path.exists():
        from lab.utils import run as shell_run

        shell_run(["docker", "compose", "-f", str(compose_path), "down", "-v"], cwd=compose_path.parent.parent)
    print({"cleaned": args.run_id})
    return 0


def cmd_retry_failures(args) -> int:
    failed = [
        item
        for item in load_json_dir(RUNS_DIR)
        if item.get("verification_status") in {"blocked-artifact", "triage-manual"}
    ]
    task_queue.enqueue_items(
        [{"advisory_id": item["advisory_id"], "system_id": item["system_id"], "priority": "retry"} for item in failed[: args.limit]]
    )
    print({"requeued": min(len(failed), args.limit)})
    return 0


def cmd_validate(args) -> int:
    errors = validators.validate_assets()
    if errors:
        print("Validation failed:")
        for error in errors:
            print(f"- {error}")
        return 1
    print("Validation passed.")
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Websafe local lab orchestrator")
    subparsers = parser.add_subparsers(dest="command", required=True)

    catalog_sync = subparsers.add_parser("catalog", help="catalog operations")
    catalog_sub = catalog_sync.add_subparsers(dest="catalog_command", required=True)
    catalog_sync_cmd = catalog_sub.add_parser("sync", help="sync environment catalog and repro map")
    catalog_sync_cmd.set_defaults(func=cmd_catalog_sync)

    compose_generate = subparsers.add_parser("compose", help="compose operations")
    compose_sub = compose_generate.add_subparsers(dest="compose_command", required=True)
    compose_generate_cmd = compose_sub.add_parser("generate", help="generate compose file for a case")
    compose_generate_cmd.add_argument("--case", required=True)
    compose_generate_cmd.add_argument("--run-id")
    compose_generate_cmd.set_defaults(func=cmd_compose_generate)

    for name, func in [
        ("provision", cmd_provision),
        ("seed", cmd_seed),
        ("baseline", cmd_baseline),
        ("attack", cmd_attack),
        ("verify", cmd_verify),
    ]:
        sub = subparsers.add_parser(name)
        sub.add_argument("--case", required=True)
        sub.add_argument("--run-id")
        sub.add_argument("--dry-run", action="store_true")
        sub.set_defaults(func=func)

    run_case = subparsers.add_parser("run-case", help="run a single advisory through the lab pipeline")
    run_case.add_argument("--case", required=True)
    run_case.add_argument("--run-id")
    run_case.add_argument("--dry-run", action="store_true")
    run_case.set_defaults(func=cmd_run_case)

    run_system = subparsers.add_parser("run-system", help="run the first N advisories for a system")
    run_system.add_argument("--system", required=True)
    run_system.add_argument("--limit", type=int, default=5)
    run_system.add_argument("--dry-run", action="store_true")
    run_system.set_defaults(func=cmd_run_system)

    run_batch = subparsers.add_parser("run-batch", help="process repro queue or enqueue from registry")
    run_batch.add_argument("--limit", type=int, default=10)
    run_batch.add_argument("--only-hotlane", action="store_true")
    run_batch.add_argument("--from-queue", action="store_true")
    run_batch.add_argument("--dry-run", action="store_true")
    run_batch.set_defaults(func=cmd_run_batch)

    render_run = subparsers.add_parser("render-run", help="re-render a stored run")
    render_run.add_argument("--run-id", required=True)
    render_run.set_defaults(func=cmd_render_run)

    serve = subparsers.add_parser("serve-dashboard", help="serve the static dashboard locally")
    serve.add_argument("--port", type=int, default=8734)
    serve.set_defaults(func=cmd_serve_dashboard)

    cleanup = subparsers.add_parser("cleanup", help="tear down a stored run compose environment")
    cleanup.add_argument("--run-id", required=True)
    cleanup.set_defaults(func=cmd_cleanup)

    retry = subparsers.add_parser("retry-failures", help="requeue blocked or manual runs")
    retry.add_argument("--limit", type=int, default=50)
    retry.set_defaults(func=cmd_retry_failures)

    validate = subparsers.add_parser("validate", help="validate lab assets")
    validate.set_defaults(func=cmd_validate)
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())