#!/usr/bin/env python3
"""Load Piko-9b once and exercise every capability the model card claims.

Each check is recorded independently, so a failure in one modality does not
hide the results of the others.  Output is a JSON record suitable for pasting
into the audit report; nothing here is scored by hand.

Usage
-----
python scripts/validate_inference.py --model <path-or-repo-id> \
    --output reports/inference_validation.json
"""

from __future__ import annotations

import argparse
import json
import platform
import sys
import time
import traceback
from collections.abc import Callable
from pathlib import Path
from typing import Any


def build_ocr_image(path: Path) -> None:
    """Render a deterministic synthetic receipt.  No network, no licensing risk."""
    from PIL import Image, ImageDraw

    image = Image.new("RGB", (520, 300), "white")
    draw = ImageDraw.Draw(image)
    lines = [
        "NORTHGATE HARDWARE",
        "144 Mill Road",
        "",
        "Date: 2026-03-14",
        "Invoice: 40817",
        "",
        "Hex bolts M6      12.40",
        "Wood glue          6.25",
        "Sandpaper pack     4.10",
        "",
        "TOTAL             22.75",
    ]
    y = 18
    for line in lines:
        draw.text((24, y), line, fill="black")
        y += 24
    image.save(path)


def build_chart_image(path: Path) -> None:
    """Render a deterministic bar chart with labelled values."""
    from PIL import Image, ImageDraw

    image = Image.new("RGB", (460, 300), "white")
    draw = ImageDraw.Draw(image)
    bars = [("Q1", 40), ("Q2", 95), ("Q3", 60), ("Q4", 130)]
    base_y = 250
    for index, (label, value) in enumerate(bars):
        x = 60 + index * 90
        draw.rectangle([x, base_y - value, x + 50, base_y], fill="black")
        draw.text((x + 12, base_y + 8), label, fill="black")
        draw.text((x + 6, base_y - value - 16), str(value), fill="black")
    draw.text((40, 12), "Units sold by quarter", fill="black")
    image.save(path)


class Validator:
    def __init__(
        self, model_path: str, dtype: str, device_map: Any, quantization: str = "none"
    ) -> None:
        self.model_path = model_path
        self.dtype = dtype
        self.device_map = device_map
        self.quantization = quantization
        self.results: list[dict[str, Any]] = []
        self.model = None
        self.processor = None
        self.tokenizer = None

    # -- harness ---------------------------------------------------------- #

    def check(self, name: str, fn: Callable[[], Any]) -> Any:
        started = time.perf_counter()
        try:
            detail = fn()
            record = {
                "check": name,
                "status": "pass",
                "seconds": round(time.perf_counter() - started, 2),
                "detail": detail,
            }
        except Exception as exc:  # noqa: BLE001 - every failure must be recorded
            record = {
                "check": name,
                "status": "fail",
                "seconds": round(time.perf_counter() - started, 2),
                "error": f"{type(exc).__name__}: {exc}",
                "traceback": traceback.format_exc(limit=4),
            }
        self.results.append(record)
        marker = "PASS" if record["status"] == "pass" else "FAIL"
        print(f"[{marker}] {name} ({record['seconds']}s)", flush=True)
        if record["status"] == "fail":
            print(f"        {record['error']}", flush=True)
        return record

    # -- loading ---------------------------------------------------------- #

    def load(self) -> dict[str, Any]:
        import torch
        from transformers import AutoConfig, AutoProcessor, AutoTokenizer

        torch_dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16}[self.dtype]
        config = AutoConfig.from_pretrained(self.model_path)

        extra: dict[str, Any] = {}
        if self.quantization == "4bit":
            from transformers import BitsAndBytesConfig

            extra["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_use_double_quant=True,
            )
        elif self.quantization == "8bit":
            from transformers import BitsAndBytesConfig

            extra["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)

        loaded_with = None
        model = None
        errors: dict[str, str] = {}
        for class_name in ("AutoModelForMultimodalLM", "AutoModelForImageTextToText"):
            try:
                import transformers

                cls = getattr(transformers, class_name)
            except AttributeError:
                errors[class_name] = "class not available in this transformers version"
                continue
            try:
                model = cls.from_pretrained(
                    self.model_path,
                    dtype=torch_dtype,
                    device_map=self.device_map,
                    **extra,
                )
                loaded_with = class_name
                break
            except Exception as exc:  # noqa: BLE001
                errors[class_name] = f"{type(exc).__name__}: {exc}"

        if model is None:
            raise RuntimeError(f"No auto class could load the model: {errors}")

        model.eval()
        self.model = model
        self.processor = AutoProcessor.from_pretrained(self.model_path)
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)

        parameters = sum(p.numel() for p in model.parameters())
        vision_parameters = 0
        if hasattr(model, "model") and hasattr(model.model, "visual"):
            vision_parameters = sum(p.numel() for p in model.model.visual.parameters())

        return {
            "loaded_with": loaded_with,
            "auto_class_errors": errors,
            "trust_remote_code_required": False,
            "architectures": config.architectures,
            "model_type": config.model_type,
            "total_parameters": parameters,
            "vision_parameters": vision_parameters,
            "language_parameters": parameters - vision_parameters,
            "device_map": str(getattr(model, "hf_device_map", self.device_map)),
            "processor_class": type(self.processor).__name__,
            "tokenizer_class": type(self.tokenizer).__name__,
        }

    # -- generation helpers ----------------------------------------------- #

    def _generate(self, messages: list[dict[str, Any]], max_new_tokens: int) -> str:
        import torch

        inputs = self.processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(self.model.device)

        with torch.inference_mode():
            output = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
        prompt_length = inputs["input_ids"].shape[1]
        return self.processor.decode(output[0][prompt_length:], skip_special_tokens=True).strip()

    def text_only(self, prompt: str, max_new_tokens: int = 96) -> dict[str, Any]:
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
        text = self._generate(messages, max_new_tokens)
        return {"prompt": prompt, "response": text}

    def with_image(
        self, image_path: Path, prompt: str, max_new_tokens: int = 128
    ) -> dict[str, Any]:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "url": str(image_path)},
                    {"type": "text", "text": prompt},
                ],
            }
        ]
        text = self._generate(messages, max_new_tokens)
        return {"image": image_path.name, "prompt": prompt, "response": text}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16"])
    parser.add_argument("--device-map", default="auto")
    parser.add_argument(
        "--quantization",
        default="none",
        choices=["none", "4bit", "8bit"],
        help="CPU offload corrupts this architecture; use 4bit to stay resident on one GPU.",
    )
    parser.add_argument("--assets", type=Path, default=Path("evaluation/prompts/assets"))
    parser.add_argument("--output", type=Path, default=Path("reports/inference_validation.json"))
    args = parser.parse_args()

    import torch
    import transformers

    args.assets.mkdir(parents=True, exist_ok=True)
    ocr_image = args.assets / "synthetic_receipt.png"
    chart_image = args.assets / "synthetic_chart.png"
    build_ocr_image(ocr_image)
    build_chart_image(chart_image)

    device_map: Any = args.device_map
    if args.quantization != "none" and device_map == "auto":
        device_map = {"": 0}  # keep every module on one device
    validator = Validator(args.model, args.dtype, device_map, args.quantization)

    environment = {
        "python": platform.python_version(),
        "platform": platform.platform(),
        "torch": torch.__version__,
        "transformers": transformers.__version__,
        "cuda_available": torch.cuda.is_available(),
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "vram_bytes": torch.cuda.get_device_properties(0).total_memory
        if torch.cuda.is_available()
        else None,
        "dtype": args.dtype,
        "device_map": str(device_map),
        "quantization": args.quantization,
        "model": args.model,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }

    load_record = validator.check("load_model", validator.load)
    if load_record["status"] == "fail":
        _write(args.output, environment, validator.results)
        sys.exit("Model failed to load; remaining checks skipped.")

    validator.check(
        "text_only_generation",
        lambda: validator.text_only("Write a Python function that reverses a string."),
    )
    validator.check(
        "text_only_identity",
        lambda: validator.text_only("What model are you? Answer in one short sentence.", 48),
    )
    validator.check(
        "text_only_reasoning",
        lambda: validator.text_only(
            "A shop sells pens at 3 for $2. How much do 12 pens cost? Answer with the number only.",
            48,
        ),
    )
    validator.check(
        "image_ocr",
        lambda: validator.with_image(
            ocr_image, "Read this receipt. Give the merchant name and the total."
        ),
    )
    validator.check(
        "image_document_json",
        lambda: validator.with_image(
            ocr_image,
            'Return only JSON: {"merchant": str, "date": "YYYY-MM-DD", "total": float}',
        ),
    )
    validator.check(
        "image_chart",
        lambda: validator.with_image(chart_image, "Which quarter is highest, and what value?"),
    )
    validator.check(
        "image_caption",
        lambda: validator.with_image(chart_image, "Describe this image in one sentence."),
    )

    def multi_turn() -> dict[str, Any]:
        messages = [
            {"role": "user", "content": [{"type": "text", "text": "My favourite number is 47."}]},
            {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]},
            {
                "role": "user",
                "content": [{"type": "text", "text": "Double my favourite number. Number only."}],
            },
        ]
        return {"response": validator._generate(messages, 32)}

    validator.check("multi_turn_conversation", multi_turn)

    def determinism() -> dict[str, Any]:
        first = validator.text_only("Name three primary colours.", 32)["response"]
        second = validator.text_only("Name three primary colours.", 32)["response"]
        return {"identical": first == second, "first": first, "second": second}

    validator.check("greedy_determinism", determinism)

    def batch() -> dict[str, Any]:
        import torch

        prompts = ["Capital of Japan?", "2 + 2 = ?"]
        texts = [
            validator.processor.apply_chat_template(
                [{"role": "user", "content": [{"type": "text", "text": p}]}],
                add_generation_prompt=True,
                tokenize=False,
            )
            for p in prompts
        ]
        inputs = validator.processor(text=texts, return_tensors="pt", padding=True).to(
            validator.model.device
        )
        with torch.inference_mode():
            output = validator.model.generate(**inputs, max_new_tokens=24, do_sample=False)
        decoded = [
            validator.processor.decode(
                output[i][inputs["input_ids"].shape[1] :], skip_special_tokens=True
            ).strip()
            for i in range(len(prompts))
        ]
        return {"prompts": prompts, "responses": decoded}

    validator.check("batch_inference", batch)

    def long_context() -> dict[str, Any]:
        needle = "The maintenance code for the north pump is QF-8812."
        filler = "Routine log entry: all systems nominal. " * 900
        prompt = f"{filler}\n{needle}\n{filler}\n\nWhat is the maintenance code for the north pump?"
        tokens = len(validator.tokenizer(prompt)["input_ids"])
        response = validator.text_only(prompt, 32)["response"]
        return {
            "prompt_tokens": tokens,
            "response": response,
            "contains_needle": "QF-8812" in response,
        }

    validator.check("long_context_retrieval", long_context)

    _write(args.output, environment, validator.results)
    passed = sum(1 for r in validator.results if r["status"] == "pass")
    print(f"\n{passed}/{len(validator.results)} checks passed -> {args.output}")


def _write(output: Path, environment: dict[str, Any], results: list[dict[str, Any]]) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "environment": environment,
        "summary": {
            "total": len(results),
            "passed": sum(1 for r in results if r["status"] == "pass"),
            "failed": sum(1 for r in results if r["status"] == "fail"),
        },
        "results": results,
    }
    output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()