from __future__ import annotations

import argparse
from pathlib import Path
from typing import Dict, Optional

from .cli_api import _load_system, _load_trajectory
from .cli_args import REGISTRY, add_shared_args
from .cli_builders import CLI_TO_PLAN, PLAN_BUILDERS
from .cli_config import _default_out, _load_config, _normalize_system_spec, _normalize_traj_spec, example_config
from .cli_output import _print_summary, _save_output, _summary_from_output
from .cli_parse import _load_system_from_args, _load_traj_from_args
from .cli_specs import SPEC_BUILDERS


def run_config(config_path: str, dry_run: bool = False) -> None:
    cfg = _load_config(config_path)
    system_spec = _normalize_system_spec(cfg.get("system") or cfg.get("topology"))
    traj_spec = _normalize_traj_spec(cfg.get("trajectory") or cfg.get("traj"))
    system = _load_system(system_spec)
    default_device = cfg.get("device", "auto")
    default_chunk = cfg.get("chunk_frames")
    output_dir = cfg.get("output_dir", ".")
    if not dry_run:
        Path(output_dir).mkdir(parents=True, exist_ok=True)

    analyses = cfg.get("analyses")
    if not analyses:
        raise ValueError("config.analyses is required")

    used_names: Dict[str, int] = {}
    for item in analyses:
        name = item.get("name")
        if name not in PLAN_BUILDERS:
            alt_name = name.replace("-", "_") if isinstance(name, str) else name
            if alt_name not in PLAN_BUILDERS:
                raise ValueError(f"unknown analysis name: {name}")
            name = alt_name
        out_path = item.get("out") or _default_out(name, output_dir, used_names)
        if dry_run:
            print(f"{name} -> {out_path}")
            continue
        plan = PLAN_BUILDERS[name](system, item)
        traj = _load_trajectory(traj_spec, system)
        device = item.get("device", default_device)
        chunk = item.get("chunk_frames", default_chunk)
        output = plan.run(traj, system, chunk_frames=chunk, device=device)
        _save_output(out_path, output)
        print(f"{name}: wrote {out_path}")


def list_plans() -> None:
    for name in sorted(CLI_TO_PLAN.keys()):
        print(name)


def build_plan_from_args(args: argparse.Namespace, system):
    plan_name = CLI_TO_PLAN[args.analysis]
    spec = SPEC_BUILDERS[args.analysis](args, system)
    return PLAN_BUILDERS[plan_name](system, spec)


def run_single_analysis(args: argparse.Namespace) -> None:
    system = _load_system_from_args(args)
    traj = _load_traj_from_args(args, system)
    plan = build_plan_from_args(args, system)
    output = plan.run(traj, system, chunk_frames=args.chunk_frames, device=args.device)
    out_path = Path(args.out or f"{args.analysis}.npz")
    _save_output(str(out_path), output)
    if args.print_summary:
        summary = _summary_from_output(output, args.analysis, out_path)
        _print_summary(summary, args.summary_format)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="warp-md")
    sub = parser.add_subparsers(dest="cmd", required=True)

    run = sub.add_parser("run", help="run analyses from a JSON/YAML config")
    run.add_argument("config", help="path to config.json|yaml")
    run.add_argument("--dry-run", action="store_true", help="validate and show outputs")

    sub.add_parser("list-plans", help="list available analysis names")
    sub.add_parser("example", help="print example config")

    for name, setup in REGISTRY.items():
        help_text = f"Run {name} analysis"
        analysis = sub.add_parser(name, help=help_text, description=help_text)
        add_shared_args(analysis)
        setup(analysis)

    return parser


def main(argv: Optional[list[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.cmd == "run":
        run_config(args.config, dry_run=args.dry_run)
        return 0
    if args.cmd == "list-plans":
        list_plans()
        return 0
    if args.cmd == "example":
        example_config()
        return 0
    if args.cmd in REGISTRY:
        args.analysis = args.cmd
        run_single_analysis(args)
        return 0

    return 1
