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

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


def _default_repo() -> Path:
    return Path(__file__).resolve().parents[1]


def _default_poscar(repo: Path) -> Path:
    return repo / "workdir" / "srti3-user-input" / "POSCAR"


def _build_parser() -> argparse.ArgumentParser:
    repo = _default_repo()
    p = argparse.ArgumentParser(
        description="Run a single macer phonopy sscha/qscaild case (default: SrTiO3 300K, 3x3x3)."
    )
    p.add_argument("--repo", type=Path, default=repo, help="macer repo root")
    p.add_argument("-p", "--poscar", type=Path, default=_default_poscar(repo), help="input POSCAR")
    p.add_argument("-T", "--temperature", type=float, default=300.0, help="temperature in K")
    p.add_argument("--dim", nargs=3, type=int, default=[3, 3, 3], help="supercell dimensions")
    p.add_argument("--nconf", type=int, default=10, help="configs per cycle")
    p.add_argument("--nsteps", type=int, default=30, help="max cycles")
    p.add_argument("--imaginary-freq", type=float, default=1.0, help="imaginary freq handling (THz)")
    p.add_argument("--tolerance", type=float, default=0.02, help="IFC convergence tolerance")
    p.add_argument("--pdiff", type=float, default=2.0, help="pressure tolerance (kbar)")
    p.add_argument("--grid", type=int, default=9, help="q-grid size")
    p.add_argument("--memory", type=float, default=0.3, help="history ratio")
    p.add_argument("--mixing", type=float, default=0.0, help="IFC mixing")
    p.add_argument("--ff", type=str, default="chgnet", help="MLFF name")
    p.add_argument("--device", choices=["cpu", "mps", "cuda"], default="cpu", help="compute device")
    p.add_argument("--save-every", type=int, default=5, help="save FC snapshots every N cycles")
    p.add_argument("--output-dir", type=Path, default=None, help="output directory")
    p.add_argument("--dry-run", action="store_true", help="validate setup only")
    return p


def main() -> None:
    args = _build_parser().parse_args()

    repo = args.repo.resolve()
    poscar = args.poscar.resolve()
    if not poscar.exists():
        raise FileNotFoundError(f"POSCAR not found: {poscar}")

    py = os.environ.get("MACER_PYTHON") or shutil.which("python") or sys.executable
    cmd = [
        py,
        "-m",
        "macer.cli.phonopy_main",
        "sscha",
        "-p",
        str(poscar),
        "-T",
        str(args.temperature),
        "--dim",
        str(args.dim[0]),
        str(args.dim[1]),
        str(args.dim[2]),
        "--nconf",
        str(args.nconf),
        "--nsteps",
        str(args.nsteps),
        "--imaginary-freq",
        str(args.imaginary_freq),
        "--tolerance",
        str(args.tolerance),
        "--pdiff",
        str(args.pdiff),
        "--grid",
        str(args.grid),
        "--memory",
        str(args.memory),
        "--mixing",
        str(args.mixing),
        "--ff",
        str(args.ff),
        "--device",
        str(args.device),
        "--save-every",
        str(args.save_every),
    ]
    if args.output_dir is not None:
        cmd += ["--output-dir", str(args.output_dir.resolve())]
    if args.dry_run:
        cmd.append("--dry-run")

    env = os.environ.copy()
    env["PYTHONPATH"] = f"{repo}:{env.get('PYTHONPATH', '')}".rstrip(":")

    print("Running command:")
    print(" ".join(cmd))
    subprocess.run(cmd, cwd=str(repo), env=env, check=True)


if __name__ == "__main__":
    main()
