import argparse
import json
import os
import statistics
import time

from model_config import SYSTEM_PROMPT, resolve_model
from model_runtime import resolve_generation_settings
from provenance import runtime_provenance
from run_questions import get_model_device, read_questions
from terminal_chat import generate, load_model


def mean(values):
    return statistics.mean(values) if values else None


def main():
    parser = argparse.ArgumentParser(
        description="Benchmark model load time and per-question inference latency."
    )
    parser.add_argument(
        "--questions",
        default="examples/questions.txt",
        help="Question file to benchmark.",
    )
    parser.add_argument(
        "--input-format",
        choices=["lines", "blocks"],
        default="lines",
    )
    parser.add_argument(
        "--model",
        help="Model name from models.json. Defaults to MODEL_NAME or models.json default.",
    )
    parser.add_argument("--base-model", help="Override the base Hugging Face model repo.")
    parser.add_argument("--adapter", help="Override the LoRA adapter Hugging Face repo.")
    parser.add_argument(
        "--model-mode",
        choices=["adapter", "base"],
        default=os.getenv("MODEL_MODE", "adapter").lower(),
    )
    parser.add_argument(
        "--device",
        choices=["auto", "cuda", "mps", "cpu"],
        default=os.getenv("DEVICE", "auto").lower(),
    )
    parser.add_argument("--allow-cpu", action="store_true")
    parser.add_argument("--fast", action="store_true")
    parser.add_argument("--max-new-tokens", type=int)
    parser.add_argument("--temperature", type=float)
    parser.add_argument("--top-p", type=float)
    parser.add_argument(
        "--limit",
        type=int,
        help="Only run the first N questions.",
    )
    args = parser.parse_args()

    questions = read_questions(args.questions, args.input_format)
    if args.limit is not None:
        questions = questions[: args.limit]
    if not questions:
        raise SystemExit("No questions to benchmark.")

    selected_model = resolve_model(args.model, args.base_model, args.adapter)
    generation_settings = resolve_generation_settings(
        args.max_new_tokens,
        args.temperature,
        args.top_p,
        args.fast,
    )

    load_started = time.perf_counter()
    tokenizer, model = load_model(
        selected_model,
        args.allow_cpu,
        args.model_mode,
        args.device,
    )
    load_seconds = time.perf_counter() - load_started

    results = []
    for index, question in enumerate(questions, start=1):
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": question},
        ]
        started = time.perf_counter()
        answer = generate(
            tokenizer,
            model,
            messages,
            generation_settings["max_new_tokens"],
            generation_settings["temperature"],
            generation_settings["top_p"],
        )
        seconds = time.perf_counter() - started
        output_tokens = len(tokenizer(answer, add_special_tokens=False)["input_ids"])
        results.append(
            {
                "index": index,
                "question": question,
                "answer": answer,
                "seconds": round(seconds, 3),
                "output_tokens": output_tokens,
                "tokens_per_second": round(output_tokens / seconds, 3)
                if seconds > 0
                else None,
            }
        )

    latencies = [item["seconds"] for item in results]
    token_rates = [
        item["tokens_per_second"]
        for item in results
        if item["tokens_per_second"] is not None
    ]
    report = {
        "device": get_model_device(model),
        "load_seconds": round(load_seconds, 3),
        "question_count": len(results),
        "latency_seconds": {
            "mean": round(mean(latencies), 3),
            "min": round(min(latencies), 3),
            "max": round(max(latencies), 3),
        },
        "tokens_per_second": {
            "mean": round(mean(token_rates), 3) if token_rates else None,
            "min": round(min(token_rates), 3) if token_rates else None,
            "max": round(max(token_rates), 3) if token_rates else None,
        },
        "fast": args.fast,
        "results": results,
        **runtime_provenance(
            selected_model,
            args.model_mode,
            args.device,
            generation_settings,
        ),
    }
    print(json.dumps(report, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()