"""Command line entrypoint for HALT-CoT."""

from __future__ import annotations

import argparse
import json

from .core import (
    HaltCoTConfig,
    integer_candidates,
    multiple_choice_candidates,
    yes_no_candidates,
)
from .transformers_backend import HaltCoTForCausalLM


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run HALT-CoT with a Hugging Face causal LM.")
    parser.add_argument("--model", required=True, help="Hugging Face model id or local model path.")
    parser.add_argument("--question", required=True, help="Question to answer.")
    parser.add_argument(
        "--candidate-set",
        choices=("custom", "yes-no", "multiple-choice", "integers"),
        default="custom",
        help="Candidate answer set.",
    )
    parser.add_argument(
        "--candidates",
        nargs="*",
        default=None,
        help="Custom candidate labels, e.g. --candidates Yes No or --candidates A B C D E.",
    )
    parser.add_argument("--integer-start", type=int, default=0)
    parser.add_argument("--integer-end", type=int, default=100)
    parser.add_argument("--theta", type=float, default=0.6)
    parser.add_argument("--max-steps", type=int, default=12)
    parser.add_argument("--min-steps", type=int, default=1)
    parser.add_argument("--consecutive", type=int, default=2)
    parser.add_argument("--step-max-new-tokens", type=int, default=96)
    parser.add_argument("--entropy-unit", choices=("bits", "nats"), default="bits")
    parser.add_argument("--device-map", default=None, help="Optional Transformers device_map, e.g. auto.")
    parser.add_argument("--trust-remote-code", action="store_true")
    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    candidates = _resolve_candidates(args)
    config = HaltCoTConfig(
        theta=args.theta,
        max_steps=args.max_steps,
        min_steps=args.min_steps,
        consecutive_low_entropy=args.consecutive,
        step_max_new_tokens=args.step_max_new_tokens,
        entropy_unit=args.entropy_unit,
    )

    runner = HaltCoTForCausalLM.from_pretrained(
        args.model,
        config=config,
        device_map=args.device_map,
        trust_remote_code=args.trust_remote_code,
    )
    result = runner.run(args.question, candidates)

    if args.json:
        print(
            json.dumps(
                {
                    "answer": result.answer,
                    "halted": result.halted,
                    "generated_tokens": result.generated_tokens,
                    "steps": [
                        {
                            "index": step.index,
                            "text": step.text,
                            "entropy": step.entropy,
                            "prediction": step.prediction,
                            "probabilities": step.probabilities,
                            "halted": step.halted,
                            "generated_tokens": step.generated_tokens,
                        }
                        for step in result.steps
                    ],
                },
                indent=2,
            )
        )
    else:
        print(f"Answer: {result.answer}")
        print(f"Halted: {result.halted}")
        print(f"Generated tokens: {result.generated_tokens}")
        print("\nTrace:")
        for step in result.steps:
            marker = " HALT" if step.halted else ""
            print(
                f"- Step {step.index}: H={step.entropy:.3f} "
                f"pred={step.prediction!r}{marker} | {step.text}"
            )
    return 0


def _resolve_candidates(args: argparse.Namespace):
    if args.candidate_set == "yes-no":
        return yes_no_candidates()
    if args.candidate_set == "multiple-choice":
        labels = args.candidates or ["A", "B", "C", "D", "E"]
        return multiple_choice_candidates(labels)
    if args.candidate_set == "integers":
        return integer_candidates(args.integer_start, args.integer_end)
    if not args.candidates:
        raise SystemExit("--candidates is required when --candidate-set custom")
    return args.candidates


if __name__ == "__main__":
    raise SystemExit(main())