#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import time
from datetime import datetime, timezone
from pathlib import Path


def _repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def _resolve_path(value: str) -> Path:
    path = Path(value).expanduser()
    if path.is_absolute():
        return path
    return _repo_root() / path


def _default_artifact_path() -> Path:
    today = datetime.now(timezone.utc).date().isoformat()
    return _repo_root() / "artifacts" / "verification" / today / "llama_champion_smoke.json"


def _load_llama_cpp():
    try:
        import llama_cpp
    except Exception as exc:  # pragma: no cover - import error is environment-specific
        raise SystemExit(
            "llama_cpp import failed; install llama-cpp-python==0.3.28 from the CPU wheel index first.") from exc
    return llama_cpp


def _generate_text(llm, prompt: str, *, max_tokens: int, temperature: float, top_p: float, seed: int):
    kwargs = {
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "seed": seed,
    }
    messages = [{"role": "user", "content": prompt}]
    try:
        response = llm.create_chat_completion(messages=messages, **kwargs)
        choice = response["choices"][0]
        text = (choice.get("message") or {}).get("content") or ""
        mode = "chat"
    except Exception:
        response = llm(f"{prompt}\n", echo=False, **kwargs)
        choice = response["choices"][0]
        text = choice.get("text") or ""
        mode = "completion"
    return text.strip(), response, mode


def main() -> int:
    parser = argparse.ArgumentParser(description="Run the local GGUF llama-cpp-python smoke check.")
    parser.add_argument("--model", default=os.environ.get("LLAMA_CHAMPION_MODEL"), help="Path to the local .gguf file (or set LLAMA_CHAMPION_MODEL).")
    parser.add_argument("--artifact-path", default=os.environ.get("LLAMA_CHAMPION_ARTIFACT_PATH") or str(_default_artifact_path()), help="Where to write the verification artifact JSON.")
    parser.add_argument("--prompt", default="Reply with a short phrase that includes the word champion.", help="Tiny prompt used for the smoke check.")
    parser.add_argument("--max-tokens", type=int, default=16)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top-p", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--n-ctx", type=int, default=256)
    parser.add_argument("--n-batch", type=int, default=64)
    parser.add_argument("--n-threads", type=int, default=max(1, min(8, (os.cpu_count() or 2) // 2)))
    args = parser.parse_args()

    if not args.model:
        parser.error("missing --model or LLAMA_CHAMPION_MODEL")

    model_path = _resolve_path(args.model)
    if not model_path.exists():
        parser.error(f"model not found: {model_path}")

    artifact_path = _resolve_path(args.artifact_path)
    artifact_path.parent.mkdir(parents=True, exist_ok=True)

    llama_cpp = _load_llama_cpp()
    started = time.perf_counter()
    llm = llama_cpp.Llama(
        model_path=str(model_path),
        n_ctx=args.n_ctx,
        n_batch=args.n_batch,
        n_threads=args.n_threads,
        n_gpu_layers=0,
        seed=args.seed,
        verbose=False,
    )
    loaded = time.perf_counter()
    text, response, mode = _generate_text(
        llm,
        args.prompt,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        seed=args.seed,
    )
    ended = time.perf_counter()

    if not text:
        raise SystemExit("llama.cpp smoke returned empty text")

    llama_cpp_version = getattr(llama_cpp, "__version__", "unknown")
    payload = {
        "success": True,
        "backend": "llama-cpp-python",
        "llama_cpp_version": llama_cpp_version,
        "mode": mode,
        "repo": _repo_root().name,
        "model_path": str(model_path.resolve()),
        "prompt": args.prompt,
        "response": text,
        "artifact_path": str(artifact_path.resolve()),
        "timing_s": {
            "load": round(loaded - started, 3),
            "total": round(ended - started, 3),
        },
        "created_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "usage": response.get("usage"),
    }
    artifact_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(payload, indent=2, sort_keys=True))
    return 0


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