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

import argparse
import subprocess
import sys
from pathlib import Path


def parse_parameters(path: Path) -> dict[str, str]:
    data: dict[str, str] = {}
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        data[k.strip()] = v.strip()
    return data


def map_to_macer_cli(params: dict[str, str]) -> list[str]:
    mapped = []
    if "T_K" in params:
        mapped += ["-T", params["T_K"]]
    if "nconf" in params:
        mapped += ["--nconf", params["nconf"]]
    if "nfits" in params:
        mapped += ["--nfits", params["nfits"]]
    if "memory" in params:
        mapped += ["--memory", params["memory"]]
    if "mixing" in params:
        mapped += ["--mixing", params["mixing"]]
    if "tolerance" in params:
        mapped += ["--tolerance", params["tolerance"]]
    if "pdiff" in params:
        mapped += ["--pdiff", params["pdiff"]]
    if "grid" in params:
        mapped += ["--grid", params["grid"]]
    if "imaginary_freq" in params:
        mapped += ["--imaginary-freq", params["imaginary_freq"]]
    if params.get("use_smalldisp", "False").lower() in ("true", "1", "yes"):
        mapped += ["--use-smalldisp"]
    if "use_pressure" in params:
        mapped += ["--use-pressure", params["use_pressure"]]
    if "pressure_diag" in params:
        vals = [x.strip() for x in params["pressure_diag"].split(",")]
        if len(vals) == 3:
            mapped += ["--pressure-diag", vals[0], vals[1], vals[2]]
    return mapped


def _find_poscar(case_dir: Path) -> Path | None:
    for cand in [case_dir / "POSCAR", case_dir / "POSCAR_input"]:
        if cand.exists():
            return cand
    return None


def _override_for_fast(mapped: list[str], fast: bool) -> list[str]:
    if not fast:
        return mapped

    # Remove existing conflicting options first.
    remove_keys = {"--nconf", "--nfits", "--dim", "--no-plot-bands", "--initial-isif"}
    out: list[str] = []
    i = 0
    while i < len(mapped):
        tok = mapped[i]
        if tok in remove_keys:
            i += 2
            continue
        if tok == "--pressure-diag":
            out.extend(mapped[i:i + 4])
            i += 4
            continue
        out.append(tok)
        i += 1

    out += ["--nconf", "1", "--nfits", "1", "--dim", "1", "1", "1", "--initial-isif", "0", "--no-plot-bands"]
    return out


def run_case(case_dir: Path, mapped: list[str], ff: str, timeout_s: int, dry_run: bool) -> tuple[int, str, str]:
    poscar = _find_poscar(case_dir)
    if poscar is None:
        return 2, "", f"POSCAR not found in {case_dir}"

    cmd = [
        sys.executable,
        "-m",
        "macer.cli.phonopy_main",
        "qscaild",
        "-p",
        str(poscar),
        "--ff",
        ff,
    ] + mapped
    if dry_run and "--dry-run" not in cmd:
        cmd.append("--dry-run")

    proc = subprocess.run(
        cmd,
        text=True,
        capture_output=True,
        check=False,
        timeout=timeout_s,
    )
    return proc.returncode, proc.stdout, proc.stderr


def main() -> int:
    p = argparse.ArgumentParser(description="Run smoke tests for qscaild-master reference cases via macer qscaild.")
    p.add_argument("--root", default="qscaild-master/test_dir", help="Reference root containing case directories.")
    p.add_argument("--ff", default="mattersim", help="Force field to use for smoke runs.")
    p.add_argument("--timeout", type=int, default=300, help="Per-case timeout in seconds.")
    p.add_argument("--dry-run", action="store_true", help="Use qscaild --dry-run mode only.")
    p.add_argument("--fast", action="store_true", help="Override to minimal runtime settings.")
    args = p.parse_args()

    root = Path(args.root).resolve()
    param_files = sorted(root.rglob("parameters"))
    if not param_files:
        print(f"No parameters files found under: {root}")
        return 1

    summary: list[tuple[str, int, str]] = []
    for f in param_files:
        case_dir = f.parent
        params = parse_parameters(f)
        mapped = map_to_macer_cli(params)
        mapped = _override_for_fast(mapped, args.fast)
        rc, out, err = run_case(case_dir, mapped, ff=args.ff, timeout_s=args.timeout, dry_run=args.dry_run)
        rel = str(f.relative_to(root))
        if rc == 0:
            msg = "OK"
        else:
            last = (err.strip().splitlines()[-1] if err.strip() else out.strip().splitlines()[-1] if out.strip() else "unknown error")
            msg = f"FAIL: {last}"
        summary.append((rel, rc, msg))
        print(f"[{rel}] {msg}")

    failed = [x for x in summary if x[1] != 0]
    print("")
    print("Summary:")
    print(f"  total={len(summary)} passed={len(summary)-len(failed)} failed={len(failed)}")
    if failed:
        for rel, _, msg in failed:
            print(f"  - {rel}: {msg}")
        return 1
    return 0


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