#!/usr/bin/env python3
"""Validate the self-contained Bina CoreML packages plus native host assets."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import coremltools as ct
import numpy as np


EOS_TOKEN_ID = 2


def load_raw(path: Path, shape: list[int], dtype: str) -> np.ndarray:
    array = np.fromfile(path, dtype=dtype)
    expected = int(np.prod(shape))
    if array.size != expected:
        raise ValueError(f"{path}: expected {expected} scalars, got {array.size}")
    return array.reshape(shape)


def contiguous_prefix(left: list[int], right: list[int]) -> int:
    for index, (actual, expected) in enumerate(zip(left, right)):
        if actual != expected:
            return index
    return min(len(left), len(right))


def through_eos(tokens: list[int]) -> list[int]:
    if EOS_TOKEN_ID in tokens:
        return tokens[: tokens.index(EOS_TOKEN_ID) + 1]
    return tokens


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--assets-dir", type=Path, required=True)
    parser.add_argument("--vision-package", type=Path, required=True)
    parser.add_argument("--prefill-package", type=Path, required=True)
    parser.add_argument("--decode-package", type=Path, required=True)
    parser.add_argument("--pixel-values", type=Path, required=True)
    parser.add_argument("--expected-receipt", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--max-new-tokens", type=int, default=64)
    args = parser.parse_args()

    started = time.perf_counter()
    assets = args.assets_dir.expanduser().resolve()
    constants = json.loads((assets / "surya_native_constants.json").read_text())
    shapes = constants["shapes"]
    files = constants["files"]
    source_dtype = "<f4" if constants["dtype"].startswith("float32") else "<f2"

    vision = ct.models.MLModel(str(args.vision_package.resolve()), compute_units=ct.ComputeUnit.CPU_ONLY)
    prefill = ct.models.MLModel(str(args.prefill_package.resolve()), compute_units=ct.ComputeUnit.CPU_ONLY)
    decode = ct.models.MLModel(str(args.decode_package.resolve()), compute_units=ct.ComputeUnit.CPU_ONLY)

    pixel_values = load_raw(args.pixel_values.resolve(), shapes["pixel_values"], "<f4")
    image_embeds = vision.predict({"pixel_values": pixel_values})["image_embeds"]
    inputs_embeds = load_raw(
        assets / files["prefill_text_embeds_base"],
        shapes["prefill_text_embeds_base"],
        source_dtype,
    ).astype(np.float32)
    inputs_embeds[:, constants["image_token_indices"], :] = image_embeds.reshape(
        1, image_embeds.shape[0], image_embeds.shape[1]
    )
    prefill_cos = load_raw(assets / files["prefill_cos"], shapes["prefill_cos"], source_dtype)
    prefill_sin = load_raw(assets / files["prefill_sin"], shapes["prefill_sin"], source_dtype)
    prefill_outputs = prefill.predict(
        {
            "inputs_embeds": inputs_embeds.astype(np.float32),
            "cos": prefill_cos.astype(np.float32),
            "sin": prefill_sin.astype(np.float32),
        }
    )

    full_keys = [np.asarray(prefill_outputs[f"full_key_{i}"], dtype=np.float32).copy() for i in range(6)]
    full_values = [np.asarray(prefill_outputs[f"full_value_{i}"], dtype=np.float32).copy() for i in range(6)]
    conv_states = [np.asarray(prefill_outputs[f"conv_state_{i}"], dtype=np.float32).copy() for i in range(18)]
    recurrent_states = [
        np.asarray(prefill_outputs[f"recurrent_state_{i}"], dtype=np.float32).copy() for i in range(18)
    ]
    token_embedding = load_raw(
        assets / files["token_embedding"], shapes["token_embedding"], source_dtype
    ).astype(np.float32)
    decode_cos = load_raw(assets / files["decode_cos"], shapes["decode_cos"], source_dtype)
    decode_sin = load_raw(assets / files["decode_sin"], shapes["decode_sin"], source_dtype)
    max_cache_length = shapes["decode_cos"][1]
    cache_length = len(constants["input_ids"])
    current = int(prefill_outputs["logits"][:, -1, :].argmax(axis=-1)[0])
    tokens = [current]

    while len(tokens) < args.max_new_tokens and current != EOS_TOKEN_ID:
        attention_mask = np.full((1, 1, 1, max_cache_length + 1), np.finfo(np.float32).min)
        attention_mask[..., :cache_length] = 0
        attention_mask[..., -1] = 0
        feed = {
            "inputs_embeds": token_embedding[current].reshape(1, 1, -1),
            "cos": decode_cos[:, cache_length : cache_length + 1, :].astype(np.float32),
            "sin": decode_sin[:, cache_length : cache_length + 1, :].astype(np.float32),
            "attention_mask": attention_mask,
        }
        for index, value in enumerate(full_keys):
            feed[f"full_key_{index}"] = value
        for index, value in enumerate(full_values):
            feed[f"full_value_{index}"] = value
        for index, value in enumerate(conv_states):
            feed[f"conv_state_{index}"] = value
        for index, value in enumerate(recurrent_states):
            feed[f"recurrent_state_{index}"] = value
        outputs = decode.predict(feed)
        next_token = int(outputs["logits"][:, -1, :].argmax(axis=-1)[0])
        for index in range(6):
            full_keys[index][:, :, cache_length : cache_length + 1, :] = outputs[
                f"new_full_key_{index}"
            ]
            full_values[index][:, :, cache_length : cache_length + 1, :] = outputs[
                f"new_full_value_{index}"
            ]
        conv_states = [
            np.asarray(outputs[f"new_conv_state_{i}"], dtype=np.float32).copy() for i in range(18)
        ]
        recurrent_states = [
            np.asarray(outputs[f"new_recurrent_state_{i}"], dtype=np.float32).copy() for i in range(18)
        ]
        cache_length += 1
        current = next_token
        tokens.append(current)

    id_to_token = json.loads((assets / "id_to_token.json").read_text())
    special_ids = set(constants["special_token_ids"])
    text = "".join(id_to_token[token] for token in tokens if token not in special_ids)
    receipt = {
        "prompt_tokens": len(constants["input_ids"]),
        "mrope_position_delta": constants["mrope_position_delta"],
        "host_asset_dtype": constants["dtype"],
        "tokens": tokens,
        "text": text,
        "stop_reason": "eos" if tokens[-1] == EOS_TOKEN_ID else "max_new_tokens",
        "seconds": time.perf_counter() - started,
    }
    if args.expected_receipt:
        expected_source = json.loads(args.expected_receipt.read_text())
        expected = through_eos(expected_source.get("native_tokens", expected_source.get("coreml_tokens", [])))
        receipt.update(
            {
                "expected_tokens": expected,
                "matched_prefix_tokens": contiguous_prefix(tokens, expected),
                "token_exact": tokens == expected,
            }
        )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(receipt, ensure_ascii=False, indent=2) + "\n")
    print(json.dumps(receipt, ensure_ascii=False, indent=2))
    if receipt.get("token_exact") is False or receipt["stop_reason"] != "eos":
        raise SystemExit(1)


if __name__ == "__main__":
    main()