from __future__ import annotations

import argparse
import json
import resource
import statistics
import time
import unicodedata
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image


ROOT = Path(__file__).resolve().parent
MODEL_DIR = ROOT / "model"
REPORT_DIR = ROOT / "reports"
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD = np.array([0.229, 0.224, 0.225], np.float32)
PAST = [f"past_k{i}" for i in range(6)] + [f"past_v{i}" for i in range(6)]


def milliseconds(operation):
    started = time.perf_counter()
    result = operation()
    return result, (time.perf_counter() - started) * 1000


def peak_rss_bytes() -> int:
    # macOS reports bytes; Linux reports KiB.
    value = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return int(value if __import__("platform").system() == "Darwin" else value * 1024)


def preprocess(image: Image.Image) -> np.ndarray:
    resized = image.convert("RGB").resize((224, 224), Image.Resampling.BICUBIC)
    values = (np.asarray(resized, np.float32) / 255.0 - MEAN) / STD
    return values.transpose(2, 0, 1)[None]


class Vocab:
    def __init__(self, path: Path):
        charset = json.loads(path.read_text(encoding="utf-8"))
        self.characters = ["", "", "", "", *charset]
        self.content_ids = {
            index + 4
            for index, character in enumerate(charset)
            if len(character) == 1
            and character not in "ーｰ〜~"
            and unicodedata.category(character)[0] in "LN"
        }

    def decode(self, tokens: list[int]) -> str:
        return "".join(self.characters[token] for token in tokens if token >= 4)


def choose_token(
    logits: np.ndarray, sequence: list[int], tokens: list[int], content_ids: set[int]
) -> int:
    logits = logits.astype(np.float64)
    for token in set(sequence):
        logits[token] = logits[token] * 1.2 if logits[token] < 0 else logits[token] / 1.2
    if tokens and tokens[-1] in content_ids:
        last = tokens[-1]
        run = 0
        for token in reversed(tokens):
            if token != last:
                break
            run += 1
        if run >= 12:
            logits[last] = -np.inf
    return int(np.argmax(logits))


def create_session(path: Path) -> ort.InferenceSession:
    options = ort.SessionOptions()
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    return ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"])


def recognize(vision, prefill, step, vocab: Vocab, pixels: np.ndarray) -> dict:
    vision_result, vision_ms = milliseconds(
        lambda: vision.run(["vision_embeds"], {"pixel_values": pixels})[0]
    )
    prefill_result, prefill_ms = milliseconds(
        lambda: prefill.run(
            None,
            {
                "vision_embeds": vision_result,
                "input_ids": np.array([[1]], np.int64),
            },
        )
    )
    logits = prefill_result[0][0, -1]
    present = prefill_result[1:]
    sequence = [1]
    tokens: list[int] = []
    position = vision_result.shape[1] + 1
    decode_started = time.perf_counter()
    step_ms = 0.0
    for _ in range(128):
        token = choose_token(logits, sequence, tokens, vocab.content_ids)
        if token == 2:
            break
        tokens.append(token)
        sequence.append(token)
        if len(tokens) >= 128:
            break
        feeds = {
            "input_ids": np.array([[token]], np.int64),
            "position_ids": np.array([[position]], np.int64),
        }
        feeds.update({name: value for name, value in zip(PAST, present)})
        step_result, elapsed = milliseconds(lambda: step.run(None, feeds))
        step_ms += elapsed
        logits = step_result[0][0, -1]
        present = step_result[1:]
        position += 1
    decode_ms = (time.perf_counter() - decode_started) * 1000
    return {
        "vision_ms": vision_ms,
        "prefill_ms": prefill_ms,
        "step_inference_ms": step_ms,
        "decode_loop_ms": decode_ms,
        "total_model_ms": vision_ms + prefill_ms + decode_ms,
        "tokens": len(tokens),
        "text": vocab.decode(tokens),
    }


def median_runs(runs: list[dict]) -> dict:
    keys = (
        "vision_ms",
        "prefill_ms",
        "step_inference_ms",
        "decode_loop_ms",
        "total_model_ms",
    )
    return {key: statistics.median(run[key] for run in runs) for key in keys}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--variant", choices=("121", "242"), required=True)
    parser.add_argument("--runs", type=int, default=4)
    parser.add_argument(
        "--image", default=str(ROOT / ".cache/demo/daf0244b038a-20260706.png")
    )
    parser.add_argument("--crop", default="515,5,172,300")
    args = parser.parse_args()

    crop = tuple(int(value) for value in args.crop.split(","))
    if len(crop) != 4:
        raise SystemExit("--crop must be x,y,width,height")
    source = Image.open(args.image)
    x, y, width, height = crop
    pixels = preprocess(source.crop((x, y, x + width, y + height)))
    vocab = Vocab(MODEL_DIR / "tokenizer/vocab.json")
    vision_name = "vision_int4.onnx" if args.variant == "121" else "vision_fp16.onnx"
    paths = {
        "vision": MODEL_DIR / "onnx" / vision_name,
        "prefill": MODEL_DIR / "onnx/decoder_prefill_int8.onnx",
        "step": MODEL_DIR / "onnx/decoder_step_int8.onnx",
    }
    baseline_rss = peak_rss_bytes()
    sessions = {}
    creation_ms = {}
    for name, path in paths.items():
        sessions[name], creation_ms[name] = milliseconds(lambda path=path: create_session(path))
    runs = [
        recognize(
            sessions["vision"], sessions["prefill"], sessions["step"], vocab, pixels
        )
        for _ in range(args.runs)
    ]
    report = {
        "variant": args.variant,
        "provider": sessions["vision"].get_providers(),
        "files": {name: {"path": path.name, "bytes": path.stat().st_size} for name, path in paths.items()},
        "bundle_unique_bytes": sum(path.stat().st_size for path in paths.values()),
        "session_creation_ms": creation_ms,
        "peak_rss_bytes": peak_rss_bytes(),
        "peak_rss_delta_from_start_bytes": peak_rss_bytes() - baseline_rss,
        "crop": crop,
        "runs": runs,
        "warm_median_excluding_first": median_runs(runs[1:] if len(runs) > 1 else runs),
    }
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    destination = REPORT_DIR / f"native-cpu-{args.variant}.json"
    destination.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()