"""Command-line interface for running uma-ase workflows."""

from __future__ import annotations

import json
import sys
import uuid
from pathlib import Path
from typing import Any, Dict, Iterable, Optional

from .cml_export import convert_log
from .md import build_md_options_from_args, build_md_paths, run_md_simulation
from .utils import parse_cli_args
from .workflows import (
    WorkflowPaths,
    build_output_paths,
    configure_logging,
    remove_if_exists,
    run_geometry_optimization,
    run_single_point,
    prepare_atoms_for_vibrations,
    run_vibrational_analysis,
    run_transition_state_search,
)


def _ensure_input(path: Path) -> bool:
    if not path.is_file():
        print(f"Input geometry not found: {path}", file=sys.stderr)
        return False
    return True


def _clean_outputs(paths: WorkflowPaths) -> None:
    remove_if_exists(
        tuple(
            path
            for path in (
                paths.trajectory,
                paths.log,
                paths.final_geometry,
                paths.freq_archive,
            )
            if path is not None
        )
    )


def _rollover_existing_log(log_path: Path) -> None:
    """Rename existing log file to a numbered backup before overwriting."""
    if not log_path.exists():
        return
    counter = 1
    while True:
        candidate = log_path.with_name(f"{log_path.stem}.{counter:02d}{log_path.suffix}")
        if not candidate.exists():
            log_path.rename(candidate)
            break
        counter += 1


def _clean_md_outputs(paths) -> None:
    remove_if_exists(
        (
            paths.traj,
            paths.log,
            paths.final_snapshot,
            paths.multi_xyz,
        )
    )


def _load_use_payload(use_path: Path) -> Dict[str, Any]:
    try:
        data = json.loads(use_path.read_text(encoding="utf-8"))
    except OSError as exc:
        raise ValueError(f"Unable to read .use file: {use_path}") from exc
    except json.JSONDecodeError as exc:
        raise ValueError(f"Invalid JSON in .use file: {use_path}") from exc
    if not isinstance(data, dict):
        raise ValueError("The .use file must contain a JSON object.")
    return data


def _materialize_use_input(use_path: Path, payload: Dict[str, Any]) -> Path:
    input_payload = payload.get("input") if isinstance(payload.get("input"), dict) else {}
    filename = input_payload.get("filename") if isinstance(input_payload, dict) else None
    if not filename:
        filename = f"{use_path.stem}.xyz"
    content = input_payload.get("content") if isinstance(input_payload, dict) else None
    if isinstance(content, str) and content.strip():
        target = use_path.parent / filename
        target.parent.mkdir(parents=True, exist_ok=True)
        try:
            existing = target.read_text(encoding="utf-8")
        except OSError:
            existing = None
        if existing != content:
            target.write_text(content, encoding="utf-8")
        return target
    path_value = input_payload.get("path") if isinstance(input_payload, dict) else None
    if path_value:
        return Path(str(path_value)).expanduser()
    raise ValueError("The .use file does not include input geometry content.")


def _apply_use_params(args, params: Dict[str, Any]) -> None:
    def _assign(name: str, value: Any, cast=None) -> None:
        if value is None:
            return
        if cast is not None:
            try:
                value = cast(value)
            except (TypeError, ValueError):
                pass
        setattr(args, name, value)

    run_type = params.get("run_type") or params.get("run_types") or params.get("run_type_requested")
    if run_type is not None:
        if isinstance(run_type, str):
            run_value = [item for item in run_type.split() if item]
        elif isinstance(run_type, (list, tuple)):
            run_value = list(run_type)
        else:
            run_value = [str(run_type)]
        _assign("run_type", [str(item).lower() for item in run_value])

    _assign("chg", params.get("chg") if "chg" in params else params.get("charge"), int)
    _assign("spin", params.get("spin"), int)
    _assign(
        "mlff_chk",
        params.get("checkpoint")
        if "checkpoint" in params
        else params.get("mlff_checkpoint") if "mlff_checkpoint" in params else params.get("mlff_chk"),
    )
    _assign("mlff_task", params.get("mlff_task"))
    _assign("optimizer", params.get("optimizer"))
    _assign("grad", params.get("grad"), float)
    _assign("iter", params.get("iter") if "iter" in params else params.get("iterations"), int)
    _assign("temp", params.get("temp") if "temp" in params else params.get("temperature"), float)
    _assign("press", params.get("press") if "press" in params else params.get("pressure"), float)
    _assign("ts_displacement", params.get("ts_displacement"), float)
    _assign("ts_dimer_separation", params.get("ts_dimer_separation"), float)
    _assign("ts_trial_step", params.get("ts_trial_step"), float)
    _assign("ts_max_step", params.get("ts_max_step"), float)
    _assign("ts_max_rot", params.get("ts_max_rot"), int)

    _assign("md_engine", params.get("md_engine"))
    _assign("md_steps", params.get("md_steps"), int)
    _assign("md_timestep_fs", params.get("md_timestep_fs"), float)
    _assign("md_temperature", params.get("md_temperature"), float)
    _assign("md_friction", params.get("md_friction"), float)
    _assign("md_traj_interval", params.get("md_traj_interval"), int)
    _assign("md_log_interval", params.get("md_log_interval"), int)
    _assign("md_relax_fs", params.get("md_relax_fs"), float)
    _assign("md_pressure_bar", params.get("md_pressure_bar"), float)
    _assign("md_compressibility_bar", params.get("md_compressibility_bar"), float)
    _assign("md_barostat_relax_fs", params.get("md_barostat_relax_fs"), float)
    _assign("md_nose_tdamp_fs", params.get("md_nose_tdamp_fs"), float)
    _assign("md_nose_tchain", params.get("md_nose_tchain"), int)
    _assign("md_nose_tloop", params.get("md_nose_tloop"), int)
    _assign("md_bussi_taut_fs", params.get("md_bussi_taut_fs"), float)
    _assign("md_andersen_prob", params.get("md_andersen_prob"), float)
    _assign("md_mtk_tdamp_fs", params.get("md_mtk_tdamp_fs"), float)
    _assign("md_mtk_pdamp_fs", params.get("md_mtk_pdamp_fs"), float)
    _assign("md_mtk_tchain", params.get("md_mtk_tchain"), int)
    _assign("md_mtk_pchain", params.get("md_mtk_pchain"), int)
    _assign("md_mtk_tloop", params.get("md_mtk_tloop"), int)
    _assign("md_mtk_ploop", params.get("md_mtk_ploop"), int)
    _assign("md_baoab_ttau_fs", params.get("md_baoab_ttau_fs"), float)
    _assign("md_baoab_ptau_fs", params.get("md_baoab_ptau_fs"), float)
    _assign("md_baoab_pmass", params.get("md_baoab_pmass"), float)
    _assign("md_baoab_pmass_factor", params.get("md_baoab_pmass_factor"), float)
    _assign("md_baoab_hydrostatic", params.get("md_baoab_hydrostatic"))
    _assign("md_melchionna_ttime_fs", params.get("md_melchionna_ttime_fs"), float)
    _assign("md_melchionna_ptime_fs", params.get("md_melchionna_ptime_fs"), float)
    _assign("md_melchionna_bulk_gpa", params.get("md_melchionna_bulk_gpa"), float)
    _assign("md_melchionna_mask", params.get("md_melchionna_mask"))
    _assign("md_use_pbc", params.get("md_use_pbc"))
    _assign("md_cell_a", params.get("md_cell_a"), float)
    _assign("md_cell_b", params.get("md_cell_b"), float)
    _assign("md_cell_c", params.get("md_cell_c"), float)
    _assign("md_cell_alpha", params.get("md_cell_alpha"), float)
    _assign("md_cell_beta", params.get("md_cell_beta"), float)
    _assign("md_cell_gamma", params.get("md_cell_gamma"), float)
    _assign("md_use_input_velocities", params.get("md_use_input_velocities"))
    _assign("md_velocities", params.get("md_velocities"))
    _assign("md_restart_traj", params.get("md_restart_traj"))

    cpu_only = params.get("cpu_only")
    if cpu_only is None:
        cpu_only = params.get("cpu")
    device_hint = params.get("device")
    if cpu_only is True or device_hint == "cpu":
        setattr(args, "cpu", True)

    if "chg" in params or "charge" in params:
        setattr(args, "_chg_explicit", True)
    if "spin" in params:
        setattr(args, "_spin_explicit", True)


def main(argv: Optional[Iterable[str]] = None) -> int:
    args = parse_cli_args(argv)
    args.run_id = uuid.uuid4().hex

    input_path = Path(args.input).expanduser()
    if input_path.suffix.lower() == ".use":
        try:
            use_payload = _load_use_payload(input_path)
            params = use_payload.get("params")
            if not isinstance(params, dict):
                params = {key: value for key, value in use_payload.items() if key != "input"}
            if isinstance(params, dict):
                _apply_use_params(args, params)
            input_path = _materialize_use_input(input_path, use_payload)
        except ValueError as exc:
            print(str(exc), file=sys.stderr)
            return 1
    if not _ensure_input(input_path):
        return 1
    args.input = str(input_path)

    run_sequence_raw = args.run_type
    if isinstance(run_sequence_raw, str):
        run_sequence_raw = [run_sequence_raw]
    run_sequence = [step.lower() for step in run_sequence_raw] or ["sp"]

    if "md" in run_sequence:
        if len(run_sequence) > 1:
            print(
                "MD runs cannot be combined with other run types. Use -run-type md only.",
                file=sys.stderr,
            )
            return 1
        paths = build_md_paths(input_path)
        _rollover_existing_log(paths.log)
        _clean_md_outputs(paths)
        with configure_logging(paths.log) as logger:
            try:
                md_options = build_md_options_from_args(args, logger=logger)
            except Exception as exc:
                logger.error("Invalid MD options: %s", exc)
                return 1
            status = run_md_simulation(args, logger, md_options=md_options, paths=paths)
            if status == 0:
                try:
                    cml_path = convert_log(paths.log)
                    logger.info("CML export written to %s", cml_path)
                except Exception as exc:
                    logger.warning("CML export failed for %s: %s", paths.log, exc)
        return status

    log_sequence = run_sequence.copy()

    paths = build_output_paths(input_path, log_sequence)
    _rollover_existing_log(paths.log)
    _clean_outputs(paths)

    with configure_logging(paths.log) as logger:
        context = None
        status = 0

        for step in run_sequence:
            if step == "sp":
                status, context = run_single_point(args, paths, logger)
            elif step == "geoopt":
                status, context = run_geometry_optimization(args, paths, logger)
            elif step == "freqs":
                if context is None:
                    status, context = prepare_atoms_for_vibrations(args, logger)
                    if status != 0:
                        break
                status, _ = run_vibrational_analysis(
                    context,
                    logger,
                    input_path,
                    args.temp,
                    args.press,
                    paths.freq_archive,
                    args=args,
                )
            elif step == "ts":
                status, context = run_transition_state_search(
                    args,
                    paths,
                    logger,
                    input_path,
                    context,
                )
            else:
                logger.error("Unknown run-type '%s'", step)
                status = 1

            if status != 0:
                break

        if status == 0:
            try:
                cml_path = convert_log(paths.log)
                logger.info("CML export written to %s", cml_path)
            except Exception as exc:
                logger.warning("CML export failed for %s: %s", paths.log, exc)

        return status


if __name__ == "__main__":
    sys.exit(main())
