#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11,<3.12"
# dependencies = [
#   "accelerate==1.2.1",
#   "coremltools==8.0",
#   "huggingface-hub==0.27.1",
#   "numpy==1.26.4",
#   "safetensors==0.4.5",
#   "torch==2.4.0",
#   "transformers==4.47.1",
# ]
# ///
"""Export Dolphin 3.0 Llama 3.2 as a stateful Core ML language model.

The exported ML Program keeps its key/value cache in Core ML state. It accepts
a prompt chunk for prefill and one or more continuation tokens for decode:

    inputIds:  int32 [1, 1...max_query_length]
    causalMask: fp16 [1, 1, query_length, 1...max_context_length]
    keyCache:   fp16 Core ML state
    valueCache: fp16 Core ML state
    logits:     fp16 [1, query_length, vocabulary_size]

The causal-mask final dimension is the absolute end position of the submitted
chunk. Callers must create a fresh Core ML state for each conversation.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from pathlib import Path
from typing import Any, Optional, Sequence, Tuple

import coremltools as ct
import numpy as np
import torch
from huggingface_hub import HfApi
from transformers.cache_utils import Cache
from transformers.models.llama.modeling_llama import (
    LLAMA_ATTENTION_CLASSES,
    LlamaAttention,
    LlamaConfig,
    LlamaForCausalLM,
    apply_rotary_pos_emb,
    repeat_kv,
)


DEFAULT_MODEL_ID = "dphn/Dolphin3.0-Llama3.2-3B"
DEFAULT_REVISION = "392a6f57223e7ccfe6ef4ebdb2ff101a42d57364"
DEFAULT_OUTPUT = "Dolphin3.0-Llama3.2-3B-stateful-int4.mlpackage"
TOKENIZER_METADATA_KEY = "co.huggingface.exporters.name"


class SliceUpdateKeyValueCache(Cache):
    """Fixed-size cache whose slices lower to Core ML state updates."""

    def __init__(
        self,
        shape: Tuple[int, ...],
        *,
        device: str | torch.device = "cpu",
        dtype: torch.dtype = torch.float16,
    ) -> None:
        super().__init__()
        self.past_seen_tokens = 0
        self.maximum_length = shape[-2]
        self.k_cache = torch.zeros(shape, dtype=dtype, device=device)
        self.v_cache = torch.zeros(shape, dtype=dtype, device=device)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[dict[str, Any]] = None,
        *,
        slice_indices: Optional[Tuple[int, int]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        del cache_kwargs
        if slice_indices is None:
            begin = self.past_seen_tokens
            end = begin + key_states.shape[-2]
        else:
            begin, end = slice_indices

        self.k_cache[layer_idx, :, : key_states.shape[1], begin:end, :] = key_states
        self.v_cache[layer_idx, :, : value_states.shape[1], begin:end, :] = value_states
        return (
            self.k_cache[layer_idx, :, :, :end, :],
            self.v_cache[layer_idx, :, :, :end, :],
        )

    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
        del layer_idx
        return self.past_seen_tokens

    def get_max_cache_shape(self) -> int:
        return self.maximum_length


class SliceUpdateLlamaAttention(LlamaAttention):
    """Llama SDPA attention with traceable in-place KV-cache updates."""

    @torch.no_grad()
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs: Any,
    ) -> Tuple[torch.Tensor, None, Optional[Cache]]:
        del output_attentions, use_cache, cache_position, kwargs
        batch_size, query_length, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(
            batch_size, query_length, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key_states = key_states.view(
            batch_size, query_length, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)
        value_states = value_states.view(
            batch_size, query_length, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)

        if position_embeddings is None:
            cos, sin = self.rotary_emb(value_states, position_ids)
        else:
            cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin
        )

        if attention_mask is None:
            raise ValueError("causalMask is required for stateful generation")
        if past_key_value is None:
            raise ValueError("KV cache is required for stateful generation")

        end_step = attention_mask.shape[-1]
        key_states, value_states = past_key_value.update(
            key_states,
            value_states,
            self.layer_idx,
            slice_indices=(end_step - query_length, end_step),
        )
        # Core ML Tools may materialize a state read as FP32 even when the
        # declared StateType is FP16. Keep attention operands explicitly
        # aligned with the projected query dtype.
        key_states = key_states.to(dtype=query_states.dtype)
        value_states = value_states.to(dtype=query_states.dtype)
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        attention_output = torch.nn.functional.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=attention_mask,
            dropout_p=0.0,
            is_causal=False,
        )
        attention_output = attention_output.transpose(1, 2).contiguous()
        attention_output = attention_output.view(
            batch_size, query_length, self.num_heads * self.head_dim
        )
        attention_output = self.o_proj(attention_output)
        return attention_output, None, past_key_value


class StatefulLlamaForCausalLM(torch.nn.Module):
    """Wrap a Transformers Llama model and expose its KV cache as buffers."""

    def __init__(
        self,
        model_id: str,
        revision: str,
        *,
        max_context_length: int,
        batch_size: int = 1,
        torch_dtype: torch.dtype = torch.float16,
        tiny_test: bool = False,
    ) -> None:
        super().__init__()
        LLAMA_ATTENTION_CLASSES["sdpa"] = SliceUpdateLlamaAttention
        if tiny_test:
            config = LlamaConfig(
                vocab_size=512,
                hidden_size=64,
                intermediate_size=128,
                num_hidden_layers=2,
                num_attention_heads=4,
                num_key_value_heads=2,
                max_position_embeddings=max_context_length,
                head_dim=16,
                use_cache=True,
            )
            config._attn_implementation = "sdpa"
            self.model = LlamaForCausalLM(config).to(dtype=torch_dtype)
        else:
            self.model = LlamaForCausalLM.from_pretrained(
                model_id,
                revision=revision,
                torch_dtype=torch_dtype,
                low_cpu_mem_usage=True,
                attn_implementation="sdpa",
            )
        self.model.config.use_cache = True

        config: LlamaConfig = self.model.config
        self.kv_cache_shape = (
            config.num_hidden_layers,
            batch_size,
            config.num_key_value_heads,
            max_context_length,
            config.head_dim,
        )
        self.kv_cache = SliceUpdateKeyValueCache(
            self.kv_cache_shape, dtype=torch_dtype
        )
        self.register_buffer("keyCache", self.kv_cache.k_cache)
        self.register_buffer("valueCache", self.kv_cache.v_cache)

    @torch.no_grad()
    def forward(
        self, input_ids: torch.LongTensor, causal_mask: torch.Tensor
    ) -> torch.Tensor:
        self.kv_cache.past_seen_tokens = causal_mask.shape[-1] - input_ids.shape[-1]
        return self.model(
            input_ids=input_ids,
            attention_mask=causal_mask,
            past_key_values=self.kv_cache,
            use_cache=True,
            return_dict=True,
        ).logits

    @torch.no_grad()
    def reset_cache(self) -> None:
        self.kv_cache.k_cache.zero_()
        self.kv_cache.v_cache.zero_()
        self.kv_cache.past_seen_tokens = 0


def causal_mask(
    query_length: int,
    end_step: int,
    *,
    dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
    """Return an additive mask for a chunk ending at ``end_step``."""
    if query_length < 1 or end_step < query_length:
        raise ValueError("Expected 1 <= query_length <= end_step")
    past_length = end_step - query_length
    columns = torch.arange(end_step).view(1, end_step)
    rows = past_length + torch.arange(query_length).view(query_length, 1)
    allowed = columns <= rows
    minimum = torch.finfo(dtype).min
    mask = torch.where(
        allowed,
        torch.zeros((), dtype=dtype),
        torch.full((), minimum, dtype=dtype),
    )
    return mask.view(1, 1, query_length, end_step)


@torch.no_grad()
def verify_torch_cache(model: StatefulLlamaForCausalLM) -> dict[str, float]:
    """Check cached decode against a one-pass full-prefix calculation."""
    vocab_size = model.model.config.vocab_size
    prompt = torch.tensor([[128000, 128257, 882, 198]], dtype=torch.int32)
    prompt = torch.remainder(prompt, vocab_size)
    continuation = torch.tensor([[42 % vocab_size]], dtype=torch.int32)

    model.reset_cache()
    prompt_logits = model(prompt, causal_mask(prompt.shape[-1], prompt.shape[-1]))
    cached_logits = model(
        continuation,
        causal_mask(continuation.shape[-1], prompt.shape[-1] + continuation.shape[-1]),
    )

    model.reset_cache()
    full_input = torch.cat((prompt, continuation), dim=-1)
    full_logits = model(
        full_input, causal_mask(full_input.shape[-1], full_input.shape[-1])
    )

    cached_last = cached_logits[:, -1, :].float()
    full_last = full_logits[:, -1, :].float()
    maximum_error = float(torch.max(torch.abs(cached_last - full_last)).item())
    mean_error = float(torch.mean(torch.abs(cached_last - full_last)).item())
    if not torch.allclose(cached_last, full_last, rtol=5e-3, atol=5e-3):
        raise RuntimeError(
            "KV-cache parity failed: "
            f"max_abs_error={maximum_error}, mean_abs_error={mean_error}"
        )
    if not torch.isfinite(prompt_logits).all():
        raise RuntimeError("Prompt logits contain non-finite values")
    return {"max_abs_error": maximum_error, "mean_abs_error": mean_error}


def package_inventory(package_path: Path) -> dict[str, Any]:
    files = []
    total_bytes = 0
    for path in sorted(item for item in package_path.rglob("*") if item.is_file()):
        digest = hashlib.sha256()
        with path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
                digest.update(chunk)
        size = path.stat().st_size
        total_bytes += size
        files.append(
            {
                "path": str(path.relative_to(package_path)),
                "bytes": size,
                "sha256": digest.hexdigest(),
            }
        )
    return {"bytes": total_bytes, "files": files}


def export_model(args: argparse.Namespace) -> dict[str, Any]:
    output_path = Path(args.output).resolve()
    if output_path.exists():
        if not args.overwrite:
            raise FileExistsError(f"Output already exists: {output_path}")
        shutil.rmtree(output_path)

    torch.manual_seed(0)
    model = StatefulLlamaForCausalLM(
        args.model_id,
        args.revision,
        max_context_length=args.max_context_length,
        tiny_test=args.tiny_test,
    ).eval()
    parity = verify_torch_cache(model)
    model.reset_cache()

    trace_ids = torch.tensor([[128000, 128257]], dtype=torch.int32)
    trace_ids = torch.remainder(trace_ids, model.model.config.vocab_size)
    trace_mask = causal_mask(trace_ids.shape[-1], trace_ids.shape[-1])
    traced = torch.jit.trace(
        model, (trace_ids, trace_mask), check_trace=False, strict=False
    )
    kv_cache_shape = model.kv_cache_shape
    del model

    query_length = ct.RangeDim(
        lower_bound=1,
        upper_bound=args.max_query_length,
        default=1,
        symbol="query_length",
    )
    end_step = ct.RangeDim(
        lower_bound=1,
        upper_bound=args.max_context_length,
        default=1,
        symbol="end_step",
    )
    inputs = [
        ct.TensorType(
            shape=(1, query_length), dtype=np.int32, name="inputIds"
        ),
        ct.TensorType(
            shape=(1, 1, query_length, end_step),
            dtype=np.float16,
            name="causalMask",
        ),
    ]
    states = [
        ct.StateType(
            wrapped_type=ct.TensorType(shape=kv_cache_shape, dtype=np.float16),
            name="keyCache",
        ),
        ct.StateType(
            wrapped_type=ct.TensorType(shape=kv_cache_shape, dtype=np.float16),
            name="valueCache",
        ),
    ]

    converted = ct.convert(
        traced,
        inputs=inputs,
        outputs=[ct.TensorType(dtype=np.float16, name="logits")],
        states=states,
        minimum_deployment_target=ct.target.iOS18,
        compute_units=ct.ComputeUnit.CPU_AND_GPU,
        skip_model_load=True,
    )
    del traced

    if args.quantize == "int4":
        op_config = ct.optimize.coreml.OpLinearQuantizerConfig(
            mode="linear_symmetric",
            dtype="int4",
            granularity="per_block",
            block_size=32,
        )
        optimization = ct.optimize.coreml.OptimizationConfig(
            global_config=op_config
        )
        converted = ct.optimize.coreml.linear_quantize_weights(
            converted, config=optimization
        )

    metadata = {
        TOKENIZER_METADATA_KEY: args.tokenizer_repo,
        "com.ales27pm.dolphin.source_model": args.model_id,
        "com.ales27pm.dolphin.source_revision": args.revision,
        "com.ales27pm.dolphin.max_context_length": str(args.max_context_length),
        "com.ales27pm.dolphin.max_query_length": str(args.max_query_length),
        "com.ales27pm.dolphin.stop_token_ids": "128256,128001,128008,128009",
        "com.ales27pm.dolphin.cache": "stateful-key-value",
        "com.ales27pm.dolphin.quantization": args.quantize,
    }
    converted._spec.description.metadata.author = "ales27pm; source model by dphn"
    converted._spec.description.metadata.license = "Llama 3.2 Community License"
    converted._spec.description.metadata.shortDescription = (
        "Stateful Core ML conversion of Dolphin 3.0 Llama 3.2 3B"
    )
    converted._spec.description.metadata.versionString = "2.0.0"
    converted._spec.description.metadata.userDefined.update(metadata)
    converted.input_description["inputIds"] = (
        "Token IDs for prompt prefill or continuation decode."
    )
    converted.input_description["causalMask"] = (
        "Additive FP16 causal mask; final dimension is the absolute end position."
    )
    converted.output_description["logits"] = (
        "FP16 next-token scores with shape [1, query_length, 128258]."
    )
    converted.save(str(output_path))

    spec = converted.get_spec()
    state_names = [state.name for state in spec.description.state]
    if state_names != ["keyCache", "valueCache"]:
        raise RuntimeError(f"Unexpected state schema: {state_names}")

    report = {
        "schema_version": 1,
        "source": {"repo_id": args.model_id, "revision": args.revision},
        "tiny_test": args.tiny_test,
        "artifact": output_path.name,
        "quantization": args.quantize,
        "minimum_deployment": {"ios": "18.0", "macos": "15.0"},
        "max_context_length": args.max_context_length,
        "max_query_length": args.max_query_length,
        "state_names": state_names,
        "torch_kv_cache_parity": parity,
        "inventory": package_inventory(output_path),
        "versions": {
            "coremltools": ct.__version__,
            "torch": torch.__version__,
        },
    }
    report_path = output_path.with_suffix(".export-report.json")
    rendered_report = json.dumps(report, indent=2, sort_keys=True)
    report_path.write_text(rendered_report + "\n")
    # Emit the evidence before any optional network mutation so an upload or
    # authorization failure cannot hide a successfully completed conversion.
    print(rendered_report, flush=True)

    if args.upload_repo:
        if args.tiny_test:
            raise RuntimeError("Refusing to upload a synthetic tiny-test artifact")
        token = os.environ.get("HF_TOKEN")
        if not token:
            raise RuntimeError("HF_TOKEN is required when --upload-repo is set")
        api = HfApi(token=token)
        api.upload_folder(
            repo_id=args.upload_repo,
            folder_path=str(output_path),
            path_in_repo=output_path.name,
            revision=args.upload_revision,
            commit_message=args.commit_message,
        )
        api.upload_file(
            repo_id=args.upload_repo,
            path_or_fileobj=str(report_path),
            path_in_repo=f"validation/{report_path.name}",
            revision=args.upload_revision,
            commit_message=f"Add export report for {output_path.name}",
        )

    return report


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    parser.add_argument("--revision", default=DEFAULT_REVISION)
    parser.add_argument(
        "--tokenizer-repo", default="ales27pm/Dolphin3.0-CoreML"
    )
    parser.add_argument("--output", default=DEFAULT_OUTPUT)
    parser.add_argument("--max-context-length", type=int, default=2048)
    parser.add_argument("--max-query-length", type=int, default=512)
    parser.add_argument("--quantize", choices=("none", "int4"), default="int4")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument(
        "--tiny-test",
        action="store_true",
        help="Use a random two-layer Llama to smoke-test the export pipeline.",
    )
    parser.add_argument("--upload-repo")
    parser.add_argument("--upload-revision", default="main")
    parser.add_argument(
        "--commit-message", default="Add functional stateful INT4 Core ML model"
    )
    return parser


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    if args.max_query_length > args.max_context_length:
        raise ValueError("max-query-length cannot exceed max-context-length")
    export_model(args)
    return 0


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