"""`zenith` command line: simulate, ask and eval."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from zenith.engines.base import EngineError, SpecError


def _fail(msg: str) -> None:
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(2)


def _simulate(args) -> None:
    from zenith.engines.registry import simulate

    try:
        result = simulate(json.loads(args.spec))
    except json.JSONDecodeError as e:
        _fail(f"spec is not valid JSON: {e}")
    except (SpecError, EngineError) as e:
        _fail(str(e))
    out = result.to_dict()
    out.pop("series")
    print(json.dumps(out, indent=2, ensure_ascii=False))
    if args.plot:
        from zenith.engines.plot import save_plot

        save_plot(result, Path(args.plot))


def _ask(args) -> None:
    from zenith.serve.pipeline import Zenith

    zenith = Zenith.load(Path(args.ckpt), Path(args.tokenizer), Path(args.config), device=resolve_device(args.device))
    r = zenith.ask(args.question, temperature=args.temperature, max_answer_tokens=args.max_answer_tokens)
    print(json.dumps({"spec": r.spec, "error": r.error, "truncated": r.truncated}, indent=2, ensure_ascii=False))
    if r.answer is not None:
        print(r.answer)
    if args.plot and r.result and r.result["series"]:
        from zenith.engines.base import Result
        from zenith.engines.plot import save_plot

        save_plot(Result(**r.result), Path(args.plot))
    if r.error:
        sys.exit(1)


def resolve_device(device: str) -> str:
    """`auto` means the GPU when torch can see one, else the CPU."""
    if device != "auto":
        return device
    import torch

    return "cuda" if torch.cuda.is_available() else "cpu"


def _eval(args) -> None:
    from zenith.evals.run_eval import run_eval, run_showcase
    from zenith.sft.generate import Example

    if args.baseline == "rules":
        from zenith.evals.baselines import RuleRouter

        predictor = RuleRouter()
    else:
        if not args.ckpt:
            _fail("--ckpt is required unless --baseline is given")
        from zenith.serve.pipeline import Zenith

        predictor = Zenith.load(Path(args.ckpt), Path(args.tokenizer), Path(args.config),
                                device=resolve_device(args.device))
    with open(args.data) as f:
        examples = [Example.from_json(l) for l in f if l.strip()]
    if args.limit:
        step = max(1, len(examples) // args.limit)
        examples = examples[::step][: args.limit]  # spread across ops, which are stored in order
    report = run_eval(predictor, examples, Path(args.out))
    if args.showcase:
        run_showcase(predictor, Path(args.out) / "showcase.md")
    print((Path(args.out) / "report.md").read_text())
    print(json.dumps(report["overall"], indent=2))


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="zenith", description="FarZenith astronomy model")
    sub = parser.add_subparsers(dest="command", required=True)
    p = sub.add_parser("simulate", help="run a physics engine on a JSON spec")
    p.add_argument("spec", help='e.g. \'{"sim":"orbit","op":"kepler","star_mass_msun":2}\'')
    p.add_argument("--plot", help="save the result's plots to this PNG path")
    p.add_argument("--device", default="auto", help="auto (GPU if available), cpu or cuda")
    p.set_defaults(func=_simulate)

    p = sub.add_parser("ask", help="ask the model a question")
    p.add_argument("question")
    p.add_argument("--ckpt", required=True)
    p.add_argument("--tokenizer", default="zenith/tokenizer/zenith-16k.json")
    p.add_argument("--config", default="configs/small.yaml")
    p.add_argument("--temperature", type=float, default=0.0)
    p.add_argument("--max-answer-tokens", type=int, default=256)
    p.add_argument("--plot", help="save the result's plots to this PNG path")
    p.add_argument("--device", default="auto", help="auto (GPU if available), cpu or cuda")
    p.set_defaults(func=_ask)

    p = sub.add_parser("eval", help="score a checkpoint or baseline on eval examples")
    p.add_argument("--ckpt")
    p.add_argument("--baseline", choices=["rules"])
    p.add_argument("--data", required=True, help="eval.jsonl from zenith.sft.generate")
    p.add_argument("--out", required=True)
    p.add_argument("--tokenizer", default="zenith/tokenizer/zenith-16k.json")
    p.add_argument("--config", default="configs/small.yaml")
    p.add_argument("--limit", type=int)
    p.add_argument("--showcase", action="store_true")
    p.add_argument("--device", default="auto", help="auto (GPU if available), cpu or cuda")
    p.set_defaults(func=_eval)
    return parser


def main(argv=None) -> None:
    args = build_parser().parse_args(argv)
    args.func(args)


if __name__ == "__main__":
    main()