#!/usr/bin/env python3
"""Generate once using local Ollama or a local Unsloth adapter. No paid API."""
from __future__ import annotations

import argparse
import hashlib
import importlib.metadata
import json
from pathlib import Path
import sys
import time
import urllib.parse
import urllib.request

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from vibegame.project import SYSTEM_PROMPT, parse_project


def file_hash(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def directory_fingerprint(path: Path) -> dict:
    files = {str(item.relative_to(path)): file_hash(item) for item in sorted(path.rglob("*"))
             if item.is_file() and item.suffix in {".json", ".safetensors", ".bin", ".model", ".gguf"}}
    return {"path": str(path.resolve()), "sha256": hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest(), "files": files}


class OllamaGenerator:
    def __init__(self, model: str, url: str = "http://127.0.0.1:11434", **kwargs):
        parsed = urllib.parse.urlparse(url)
        if parsed.scheme != "http" or parsed.hostname not in {"127.0.0.1", "localhost", "::1"} or parsed.username:
            raise ValueError("Only local, unauthenticated Ollama HTTP is allowed")
        self.url, self.model = url.rstrip("/"), model
        tags = self._post("/api/tags", None)
        matching = [item for item in tags.get("models", []) if item.get("name") in {model, model + ":latest"} or item.get("model") in {model, model + ":latest"}]
        if not matching:
            raise RuntimeError("Model is not installed in local Ollama; automatic pull is disabled")
        self.metadata = {"backend": "ollama-local", "model": model, "model_digest": matching[0]["digest"],
                         "runtime": self._post("/api/version", None), "system_prompt_sha256": hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()}

    def _post(self, endpoint: str, payload):
        body = None if payload is None else json.dumps(payload).encode()
        request = urllib.request.Request(self.url + endpoint, data=body, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=900) as response:
            return json.load(response)

    def generate(self, prompt: str, max_new_tokens: int = 3072) -> dict:
        start = time.monotonic()
        result = self._post("/api/generate", {"model": self.model, "system": SYSTEM_PROMPT, "prompt": prompt,
                         "stream": False, "options": {"temperature": 0, "seed": 42, "num_predict": max_new_tokens, "num_ctx": 4096}})
        elapsed = time.monotonic() - start
        return {"output": result["response"], "generation_seconds": elapsed, "generation_call_seconds": elapsed,
                "tokens": result.get("eval_count"), "done_reason": result.get("done_reason")}


class UnslothGenerator:
    def __init__(self, adapter: str, base_model: str | None = None, **kwargs):
        from unsloth import FastLanguageModel
        import torch
        if not torch.cuda.is_available():
            raise RuntimeError("Unsloth generation requires CUDA; local Ollama is the CPU/Mac alternative")
        adapter_path = Path(adapter).expanduser().resolve(strict=True)
        if not (adapter_path / "adapter_config.json").exists():
            raise ValueError("Expected a local saved PEFT adapter")
        config = json.loads((adapter_path / "adapter_config.json").read_text())
        declared_base = config["base_model_name_or_path"]
        manifest = json.loads((adapter_path / "run_manifest.json").read_text())
        signature = manifest["signature"]
        pinned_base, pinned_revision = signature["base_model"], signature["base_revision"]
        if base_model is not None and base_model != pinned_base:
            raise ValueError("Requested base model differs from the pinned training manifest")
        # PEFT adapter configs do not necessarily preserve the training revision.
        # Load the exact pinned quantized base first, then attach local LoRA weights.
        self.model, self.tokenizer = FastLanguageModel.from_pretrained(model_name=pinned_base, revision=pinned_revision,
                    max_seq_length=4096, dtype=torch.float16, load_in_4bit=True, trust_remote_code=False, use_exact_model_name=True)
        from bitsandbytes.nn import Linear4bit
        for module in self.model.modules():
            if isinstance(module, Linear4bit):
                module.compute_dtype = torch.float16
                module.compute_type_is_set = True
        from peft import PeftModel
        from transformers import AutoTokenizer
        self.model = PeftModel.from_pretrained(self.model, str(adapter_path), is_trainable=False)
        self.tokenizer = AutoTokenizer.from_pretrained(str(adapter_path), trust_remote_code=False)
        FastLanguageModel.for_inference(self.model)
        self.torch = torch
        self.metadata = {"backend": "unsloth", "adapter": directory_fingerprint(adapter_path), "base_model": pinned_base, "adapter_declared_base": declared_base,
                         "requested_base_model": base_model, "base_revision": pinned_revision,
                         "torch_version": torch.__version__, "cuda_version": torch.version.cuda,
                         "packages": {name: importlib.metadata.version(name) for name in ["unsloth", "transformers", "peft", "bitsandbytes"]},
                         "gpu": torch.cuda.get_device_name(), "system_prompt_sha256": hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()}
        try:
            from huggingface_hub import snapshot_download
            self.metadata["base_artifacts"] = directory_fingerprint(Path(snapshot_download(repo_id=pinned_base, revision=pinned_revision, local_files_only=True)))
        except Exception as exc:
            self.metadata["base_artifacts_error"] = type(exc).__name__ + ": " + str(exc)

    def generate(self, prompt: str, max_new_tokens: int = 3072) -> dict:
        return self.generate_many([prompt], max_new_tokens)[0]

    def generate_many(self, prompts: list[str], max_new_tokens: int = 3072) -> list[dict]:
        call_started = time.monotonic()
        if not 1 <= len(prompts) <= 32:
            raise ValueError("Free T4 inference batches must contain one to thirty-two prompts")
        texts = [self.tokenizer.apply_chat_template(
            [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}],
            tokenize=False, add_generation_prompt=True) for prompt in prompts]
        previous_padding = self.tokenizer.padding_side
        self.tokenizer.padding_side = "left"
        try:
            inputs = self.tokenizer(texts, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda")
        finally:
            self.tokenizer.padding_side = previous_padding
        length = inputs.input_ids.shape[1]
        if length + max_new_tokens > 4096:
            raise ValueError("Prompt plus generation exceeds the 4096-token bound; do not truncate evaluation prompts")
        self.torch.cuda.synchronize()
        start = time.monotonic()
        with self.torch.inference_mode():
            output = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False,
                                         use_cache=True, pad_token_id=self.tokenizer.eos_token_id)
        self.torch.cuda.synchronize()
        elapsed = time.monotonic() - start
        eos = self.model.generation_config.eos_token_id
        eos_ids = set(eos if isinstance(eos, list) else [eos, self.tokenizer.eos_token_id])
        results = []
        for index, sequence in enumerate(output[:, length:].tolist()):
            stop = next((i for i, token in enumerate(sequence) if token in eos_ids), None)
            generated = sequence if stop is None else sequence[:stop + 1]
            results.append({"output": self.tokenizer.decode(generated, skip_special_tokens=True),
                "generation_seconds": elapsed, "generation_batch_size": len(prompts),
                "latency_note": "generation_seconds times model.generate; generation_call_seconds includes chat formatting, tokenization, device transfers and decoding for the whole batch. Neither is divided by batch size.",
                "tokens": len(generated), "prompt_tokens": int(inputs.attention_mask[index].sum().item()),
                "done_reason": "stop" if stop is not None else "length"})
        call_elapsed = time.monotonic() - call_started
        for result in results:
            result["generation_call_seconds"] = call_elapsed
        return results


def add_backend_args(parser):
    parser.add_argument("--backend", choices=["ollama", "unsloth"], required=True)
    parser.add_argument("--model", help="Installed Ollama model tag")
    parser.add_argument("--ollama-url", default="http://127.0.0.1:11434")
    parser.add_argument("--adapter", help="Local Unsloth adapter directory")
    parser.add_argument("--base-model", default=None)
    parser.add_argument("--max-new-tokens", type=int, default=3072)


def make_generator(args):
    if not 1 <= args.max_new_tokens <= 3072:
        raise ValueError("Generation must use 1–3072 tokens, reserving context for instructions")
    if args.backend == "ollama":
        if not args.model:
            raise ValueError("--model is required for Ollama")
        return OllamaGenerator(args.model, args.ollama_url)
    if not args.adapter:
        raise ValueError("--adapter is required for Unsloth")
    return UnslothGenerator(args.adapter, args.base_model)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    add_backend_args(parser)
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if args.output.exists():
        parser.error("Output already exists; refusing to overwrite a generation")
    generator = make_generator(args)
    result = generator.generate(args.prompt, args.max_new_tokens)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("x") as handle:
        handle.write(result["output"])
    result["metadata"] = generator.metadata
    args.output.with_suffix(args.output.suffix + ".generation.json").write_text(json.dumps(result, indent=2) + "\n")
    parse_project(result["output"])
    print(args.output)