"""Workspace launcher for GP regression benchmark runs.

This script is called by the GUI and handles:
1. Provisioning/reusing a dedicated benchmark virtual environment.
2. Installing benchmark dependencies from requirements.txt.
3. Executing the GP regression pipeline (single or batch mode).
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path


BENCHMARK_NAME = "GP_regressor"
DEPS_STATE_FILE = ".deps_state.json"
DATA_SENTINEL = ".morehub_data.json"
RUN_META_FILE = ".run_meta.json"


def _log(message: str) -> None:
    print(message, flush=True)


def _cmd_display(cmd: list[str]) -> str:
    if os.name == "nt":
        return subprocess.list2cmdline(cmd)
    return " ".join(cmd)


def _run_checked(cmd: list[str], cwd: Path | None = None) -> None:
    _log(f"Running: {_cmd_display(cmd)}")
    subprocess.run(cmd, cwd=str(cwd) if cwd else None, check=True)


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _venv_python(venv_dir: Path) -> Path:
    if os.name == "nt":
        return venv_dir / "Scripts" / "python.exe"
    return venv_dir / "bin" / "python"


def _ensure_venv(venv_dir: Path) -> Path:
    python_path = _venv_python(venv_dir)
    if python_path.exists():
        _log(f"Benchmark virtual environment already exists: {venv_dir}")
        return python_path

    _log(f"Creating benchmark virtual environment: {venv_dir}")
    _run_checked([sys.executable, "-m", "venv", str(venv_dir)])

    if not python_path.exists():
        raise RuntimeError(f"Failed to create benchmark Python at {python_path}")
    return python_path


def _load_json(path: Path) -> dict:
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _save_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


def _ensure_dependencies(venv_python: Path, requirements_path: Path, state_path: Path) -> None:
    requirements_hash = _sha256(requirements_path)
    state = _load_json(state_path)

    if (
        state.get("requirements_sha256") == requirements_hash
        and state.get("requirements_path") == str(requirements_path)
        and state.get("venv_python") == str(venv_python)
    ):
        _log("Benchmark dependencies already satisfied. Skipping pip install.")
        return

    _log("Installing benchmark dependencies...")
    _run_checked([str(venv_python), "-m", "pip", "install", "--upgrade", "pip"])
    _run_checked([str(venv_python), "-m", "pip", "install", "-r", str(requirements_path)])

    _save_json(
        state_path,
        {
            "benchmark": BENCHMARK_NAME,
            "requirements_path": str(requirements_path),
            "requirements_sha256": requirements_hash,
            "venv_python": str(venv_python),
            "installed_at": datetime.now(timezone.utc).isoformat(),
        },
    )


def _safe_name(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip()).strip("_") or "run"


def _next_run_root(benchmark_root: Path) -> tuple[Path, int]:
    max_index = 0
    pattern = re.compile(r"^run_(\d+)$")

    for candidate in benchmark_root.iterdir():
        if not candidate.is_dir():
            continue
        match = pattern.match(candidate.name)
        if not match:
            continue
        max_index = max(max_index, int(match.group(1)))

    run_index = max_index + 1
    run_root = benchmark_root / f"run_{run_index}"
    run_root.mkdir(parents=True, exist_ok=False)
    return run_root, run_index


def _write_run_meta(run_root: Path, payload: dict) -> Path:
    path = run_root / RUN_META_FILE
    _save_json(path, payload)
    return path


def _derive_model_label(embeddings_path: Path, embeddings_root: Path) -> str:
    try:
        rel = embeddings_path.resolve().relative_to(embeddings_root.resolve())
        if rel.parts:
            return _safe_name(rel.parts[0])
    except ValueError:
        pass
    return _safe_name(embeddings_path.parent.name or "model")


def _build_gp_command(
    venv_python: Path,
    gp_main_script: Path,
    embeddings_path: Path,
    target_path: Path,
    splits_dir: Path,
    kernels: list[str],
    compute_rmse: bool,
    n_splits: int,
    train_ratio: float,
    regenerate_splits: bool,
    results_dir: Path,
) -> list[str]:
    cmd = [
        str(venv_python),
        str(gp_main_script),
        "--embeddings_path",
        str(embeddings_path),
        "--target_path",
        str(target_path),
        "--splits_dir",
        str(splits_dir),
        "--kernel_names",
        *kernels,
        "--compute_rmse",
        "1" if compute_rmse else "0",
        "--n_splits",
        str(n_splits),
        "--train_ratio",
        f"{train_ratio}",
        "--results_dir",
        str(results_dir),
    ]
    if regenerate_splits:
        cmd.append("--regenerate_splits")
    return cmd


def _write_manifest(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


def _new_case_dir(run_root: Path, model_label: str) -> Path:
    token = _safe_name(model_label)
    case_dir = run_root / token
    suffix = 2
    while case_dir.exists():
        case_dir = run_root / f"{token}_{suffix}"
        suffix += 1
    case_dir.mkdir(parents=True, exist_ok=False)
    return case_dir


def _run_single_case(
    *,
    venv_python: Path,
    gp_main_script: Path,
    gp_source_dir: Path,
    run_root: Path,
    data_root: Path,
    embeddings_path: Path,
    target_path: Path,
    kernels: list[str],
    compute_rmse: bool,
    n_splits: int,
    train_ratio: float,
    regenerate_splits: bool,
    mode: str,
    model_label: str,
    case_label: str,
) -> dict:
    case_dir = _new_case_dir(run_root, model_label)
    splits_dir = case_dir / "splits"
    results_dir = case_dir / "results"

    command = _build_gp_command(
        venv_python=venv_python,
        gp_main_script=gp_main_script,
        embeddings_path=embeddings_path,
        target_path=target_path,
        splits_dir=splits_dir,
        kernels=kernels,
        compute_rmse=compute_rmse,
        n_splits=n_splits,
        train_ratio=train_ratio,
        regenerate_splits=regenerate_splits,
        results_dir=results_dir,
    )

    exit_code = 0
    try:
        _run_checked(command, cwd=gp_source_dir)
        _log(f"[GP] Run completed successfully: {case_dir}")
    except subprocess.CalledProcessError as exc:
        exit_code = int(exc.returncode)
        _log(f"[GP] Run failed with exit code {exit_code}: {case_dir}")

    manifest = {
        "tool": "gp_regression",
        "mode": mode,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "data_root": str(data_root),
        "run_root": str(run_root),
        "model_label": model_label,
        "case_label": case_label,
        "input": {
            "embeddings_path": str(embeddings_path),
            "target_path": str(target_path),
        },
        "config": {
            "kernels": kernels,
            "compute_rmse": compute_rmse,
            "n_splits": n_splits,
            "train_ratio": train_ratio,
            "regenerate_splits": regenerate_splits,
        },
        "command": command,
        "output": {
            "case_dir": str(case_dir),
            "splits_dir": str(splits_dir),
            "results_dir": str(results_dir),
        },
        "status": "success" if exit_code == 0 else "failed",
        "exit_code": exit_code,
    }
    _write_manifest(case_dir / "manifest.json", manifest)
    _log(f"[GP] Run manifest: {case_dir / 'manifest.json'}")

    return {
        "status": manifest["status"],
        "exit_code": exit_code,
        "manifest_path": str(case_dir / "manifest.json"),
        "case_dir": str(case_dir),
        "model_label": model_label,
        "case_label": case_label,
    }


def _run_single_mode(
    args: argparse.Namespace,
    venv_python: Path,
    gp_main_script: Path,
    gp_source_dir: Path,
    run_root: Path,
    data_root: Path,
    embeddings_root: Path,
) -> dict:
    embeddings_path = Path(args.embeddings_path).resolve()
    target_path = Path(args.targets_path).resolve()
    if not embeddings_path.exists():
        raise FileNotFoundError(f"Embeddings file does not exist: {embeddings_path}")
    if not target_path.exists():
        raise FileNotFoundError(f"Targets file does not exist: {target_path}")

    try:
        embeddings_path.relative_to(embeddings_root.resolve())
    except ValueError:
        raise RuntimeError(
            "Single embeddings file must be inside <data_instance>/data/embeddings. "
            f"Provided: {embeddings_path}"
        )

    model_label = _derive_model_label(embeddings_path, embeddings_root)
    item = _run_single_case(
        venv_python=venv_python,
        gp_main_script=gp_main_script,
        gp_source_dir=gp_source_dir,
        run_root=run_root,
        data_root=data_root,
        embeddings_path=embeddings_path,
        target_path=target_path,
        kernels=args.kernel,
        compute_rmse=args.compute_rmse == "1",
        n_splits=args.n_splits,
        train_ratio=args.train_ratio,
        regenerate_splits=args.regenerate_splits == "1",
        mode="single",
        model_label=model_label,
        case_label=embeddings_path.stem,
    )

    return {
        "exit_code": item["exit_code"],
        "items": [item],
        "failures": 0 if item["exit_code"] == 0 else 1,
        "total": 1,
    }


def _run_batch_mode(
    args: argparse.Namespace,
    venv_python: Path,
    gp_main_script: Path,
    gp_source_dir: Path,
    run_root: Path,
    data_root: Path,
    embeddings_root: Path,
) -> dict:
    root_path = Path(args.root_path).resolve()
    target_path = Path(args.targets_path).resolve()

    if not root_path.exists() or not root_path.is_dir():
        raise FileNotFoundError(f"Batch root directory does not exist: {root_path}")
    if not target_path.exists():
        raise FileNotFoundError(f"Targets file does not exist: {target_path}")

    try:
        root_path.relative_to(embeddings_root.resolve())
    except ValueError:
        raise RuntimeError(
            "Batch root must be inside <data_instance>/data/embeddings. "
            f"Provided: {root_path}"
        )

    embeddings_files = sorted(p for p in root_path.rglob("*.npy") if p.is_file())
    if not embeddings_files:
        raise FileNotFoundError(f"No .npy embeddings files found under: {root_path}")

    _log(f"[GP] Batch mode: found {len(embeddings_files)} embeddings files")

    failures = 0
    items = []
    for embedding_path in embeddings_files:
        model_label = _derive_model_label(embedding_path, embeddings_root)
        case_label = embedding_path.stem

        item = _run_single_case(
            venv_python=venv_python,
            gp_main_script=gp_main_script,
            gp_source_dir=gp_source_dir,
            run_root=run_root,
            data_root=data_root,
            embeddings_path=embedding_path,
            target_path=target_path,
            kernels=args.kernel,
            compute_rmse=args.compute_rmse == "1",
            n_splits=args.n_splits,
            train_ratio=args.train_ratio,
            regenerate_splits=args.regenerate_splits == "1",
            mode="batch",
            model_label=model_label,
            case_label=case_label,
        )
        items.append(item)
        if item["exit_code"] != 0:
            failures += 1

    session_manifest = {
        "tool": "gp_regression",
        "mode": "batch",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "data_root": str(data_root),
        "run_root": str(run_root),
        "batch_root": str(root_path),
        "target_path": str(target_path),
        "summary": {
            "total": len(items),
            "failures": failures,
            "successes": len(items) - failures,
        },
        "items": items,
    }
    _write_manifest(run_root / "manifest.json", session_manifest)
    _log(f"[GP] Session manifest: {run_root / 'manifest.json'}")
    _log(f"[GP] Batch summary: failures={failures}, total={len(items)}")

    return {
        "exit_code": 1 if failures > 0 else 0,
        "items": items,
        "failures": failures,
        "total": len(items),
    }


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run GP benchmark with managed benchmark virtual environment")
    parser.add_argument("--data_root", required=True, help="Absolute data instance root path (data_<name>)")
    parser.add_argument("--mode", choices=["single", "batch"], required=True)

    parser.add_argument("--embeddings_path", default=None, help="Single mode: embeddings .npy path")
    parser.add_argument("--targets_path", default=None, help="Targets file path")
    parser.add_argument("--root_path", default=None, help="Batch mode: root folder to scan for .npy files")

    parser.add_argument("--kernel", action="append", choices=["gaussian", "laplacian", "matern52"], default=[])
    parser.add_argument("--compute_rmse", choices=["0", "1"], default="1")
    parser.add_argument("--regenerate_splits", choices=["0", "1"], default="0")
    parser.add_argument("--n_splits", type=int, default=5)
    parser.add_argument("--train_ratio", type=float, default=0.7)

    args = parser.parse_args()

    if not args.targets_path:
        parser.error("--targets_path is required")

    if args.mode == "single":
        if not args.embeddings_path:
            parser.error("single mode requires --embeddings_path")
    if args.mode == "batch" and not args.root_path:
        parser.error("batch mode requires --root_path")

    if not args.kernel:
        args.kernel = ["gaussian"]

    if args.n_splits < 1:
        parser.error("--n_splits must be >= 1")
    if not (0.0 < args.train_ratio <= 1.0):
        parser.error("--train_ratio must be in (0.0, 1.0]")

    return args


def main() -> int:
    args = _parse_args()

    data_root = Path(args.data_root).resolve()
    if not data_root.exists() or not data_root.is_dir():
        raise FileNotFoundError(f"Data root directory does not exist: {data_root}")
    if not (data_root / DATA_SENTINEL).exists():
        raise RuntimeError(f"Invalid data instance (missing {DATA_SENTINEL}): {data_root}")

    benchmark_root = data_root / "data" / "benchmarks" / BENCHMARK_NAME
    embeddings_root = data_root / "data" / "embeddings"
    venv_dir = benchmark_root / ".venv"

    gp_source_dir = Path(__file__).resolve().parent
    gp_main_script = gp_source_dir / "main.py"
    requirements_path = gp_source_dir / "requirements.txt"

    benchmark_root.mkdir(parents=True, exist_ok=True)
    embeddings_root.mkdir(parents=True, exist_ok=True)

    if not gp_main_script.exists():
        raise FileNotFoundError(f"GP main script not found: {gp_main_script}")
    if not requirements_path.exists():
        raise FileNotFoundError(f"GP requirements file not found: {requirements_path}")

    run_root, run_index = _next_run_root(benchmark_root)
    created_at = datetime.now(timezone.utc).isoformat()

    run_meta = {
        "format": "morehub-gp-run",
        "version": 1,
        "run_index": run_index,
        "mode": args.mode,
        "created_at": created_at,
        "status": "running",
        "data_root": str(data_root),
        "run_root": str(run_root),
    }
    _write_run_meta(run_root, run_meta)

    _log(f"[GP] Data root: {data_root}")
    _log(f"[GP] Run root: {run_root}")

    venv_python = _ensure_venv(venv_dir)
    _ensure_dependencies(venv_python, requirements_path, benchmark_root / DEPS_STATE_FILE)

    if args.mode == "single":
        outcome = _run_single_mode(
            args,
            venv_python,
            gp_main_script,
            gp_source_dir,
            run_root,
            data_root,
            embeddings_root,
        )
    else:
        outcome = _run_batch_mode(
            args,
            venv_python,
            gp_main_script,
            gp_source_dir,
            run_root,
            data_root,
            embeddings_root,
        )

    finished_at = datetime.now(timezone.utc).isoformat()
    run_meta.update(
        {
            "finished_at": finished_at,
            "status": "success" if outcome["exit_code"] == 0 else "failed",
            "summary": {
                "total": outcome.get("total", 0),
                "failures": outcome.get("failures", 0),
                "successes": max(0, outcome.get("total", 0) - outcome.get("failures", 0)),
            },
            "items": outcome.get("items", []),
        }
    )
    _write_run_meta(run_root, run_meta)

    return int(outcome["exit_code"])


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