from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from .client import LLMClient, LLMClientConfig
from .pipeline import run_evaluation
from .skeletons import generate_skeletons


def _load_task(path: Path) -> dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    required = {"count", "length", "seed", "output_dir"}
    missing = required - payload.keys()
    if missing:
        raise ValueError(f"task is missing fields: {', '.join(sorted(missing))}")
    return payload


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="logicbench")
    subparsers = parser.add_subparsers(dest="command", required=True)

    generate = subparsers.add_parser(
        "generate", help="generate hidden systems and answer-bearing behavior skeletons"
    )
    generate.add_argument("--count", type=int, default=3)
    generate.add_argument("--length", default="short")
    generate.add_argument("--seed", type=int, default=0)
    generate.add_argument("--entity-count", type=int, default=128)
    generate.add_argument("--simulation-steps", type=int, default=1000)
    generate.add_argument("--output", type=Path)

    run = subparsers.add_parser("run", help="run renderer and solver from a task JSON file")
    run.add_argument("task", type=Path)
    run.add_argument(
        "--render-chunk-size",
        type=int,
        help="override the task's records per renderer request (default: 60)",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if args.command == "generate":
        skeletons = generate_skeletons(
            args.count,
            args.length,
            args.seed,
            entity_count=args.entity_count,
            simulation_steps=args.simulation_steps,
        )
        lines = [json.dumps(item.to_dict(), ensure_ascii=False) for item in skeletons]
        text = "\n".join(lines) + "\n"
        if args.output:
            args.output.parent.mkdir(parents=True, exist_ok=True)
            args.output.write_text(text, encoding="utf-8")
        else:
            print(text, end="")
        return 0

    task = _load_task(args.task)
    renderer = LLMClient(LLMClientConfig.from_env("renderer"))
    solver = LLMClient(LLMClientConfig.from_env("solver"))
    render_chunk_size = (
        args.render_chunk_size
        if args.render_chunk_size is not None
        else int(task.get("render_chunk_size", 60))
    )
    summary = run_evaluation(
        renderer=renderer,
        solver=solver,
        count=int(task["count"]),
        length=task["length"],
        seed=int(task["seed"]),
        output_dir=Path(task["output_dir"]),
        entity_count=int(task.get("entity_count", 128)),
        simulation_steps=int(task.get("simulation_steps", 1000)),
        render_chunk_size=render_chunk_size,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0