#!/usr/bin/env python3
"""Convert the selected purpose-lite checkpoint to a fixed-shape Core ML package."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import shutil
import sys
from pathlib import Path
from typing import Any, Sequence

import numpy as np
import torch
import torch.nn.functional as functional

from purpose_data import LABELS, DataError, load_jsonl, write_json
from train import MAX_LENGTH, encode_fixed_shape


SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_DIR = (
    SCRIPT_DIR
    / "outputs"
    / "purpose-lite-v1-distilled-qat-mlx-4e"
    / "model"
)
DEFAULT_VALIDATION = SCRIPT_DIR / ".artifacts" / "dataset-v1" / "validation.jsonl"
DEFAULT_OUTPUT = (
    SCRIPT_DIR
    / "outputs"
    / "purpose-lite-v1-distilled-qat-mlx-4e"
    / "coreml"
    / "purpose-lite-v1-fp16.mlpackage"
)
COREMLTOOLS_VERSION = "9.0"


class FixedShapeBertForCoreML(torch.nn.Module):
    """Conversion-only BERT forward without Transformers' dynamic mask helpers."""

    def __init__(self, model: Any) -> None:
        super().__init__()
        self.model = model
        config = model.config
        self.num_heads = int(config.num_attention_heads)
        self.head_size = int(config.hidden_size) // self.num_heads
        if int(config.hidden_size) % self.num_heads:
            raise DataError("BERT hidden size must be divisible by attention heads")
        if int(config.max_position_embeddings) < MAX_LENGTH:
            raise DataError("BERT checkpoint cannot represent the fixed 128-token input")
        self.register_buffer(
            "fixed_position_ids",
            torch.arange(MAX_LENGTH, dtype=torch.int64).reshape(1, MAX_LENGTH),
            persistent=False,
        )

    def forward(
        self,
        input_ids: Any,
        attention_mask: Any,
        token_type_ids: Any,
    ) -> Any:
        bert = self.model.bert
        input_ids = input_ids.to(torch.int64)
        token_type_ids = token_type_ids.to(torch.int64)
        value = (
            bert.embeddings.word_embeddings(input_ids)
            + bert.embeddings.position_embeddings(self.fixed_position_ids)
            + bert.embeddings.token_type_embeddings(token_type_ids)
        )
        value = bert.embeddings.LayerNorm(value)
        zero = torch.zeros((), dtype=value.dtype, device=value.device)
        hidden = torch.full((), -10000.0, dtype=value.dtype, device=value.device)
        additive_mask = torch.where(
            attention_mask[:, None, None, :] != 0,
            zero,
            hidden,
        )

        for layer in bert.encoder.layer:
            self_attention = layer.attention.self

            def split_heads(projected: Any) -> Any:
                return projected.reshape(
                    1,
                    MAX_LENGTH,
                    self.num_heads,
                    self.head_size,
                ).permute(0, 2, 1, 3)

            queries = split_heads(self_attention.query(value))
            keys = split_heads(self_attention.key(value))
            values = split_heads(self_attention.value(value))
            scores = torch.matmul(queries, keys.transpose(-1, -2)) / math.sqrt(
                self.head_size
            )
            probabilities = torch.softmax(scores + additive_mask, dim=-1)
            context = torch.matmul(probabilities, values)
            context = context.permute(0, 2, 1, 3).reshape(
                1,
                MAX_LENGTH,
                -1,
            )
            attention_output = layer.attention.output.dense(context)
            value = layer.attention.output.LayerNorm(value + attention_output)
            intermediate = functional.gelu(
                layer.intermediate.dense(value),
                approximate="none",
            )
            value = layer.output.LayerNorm(
                value + layer.output.dense(intermediate)
            )

        pooled = torch.tanh(bert.pooler.dense(value[:, 0]))
        return self.model.classifier(pooled)


def _checkpoint_config(model_dir: Path) -> dict[str, Any]:
    config_path = model_dir / "config.json"
    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise DataError(f"{config_path}: cannot load model config: {exc}") from exc
    configured_labels = [
        config.get("id2label", {}).get(
            str(index),
            config.get("id2label", {}).get(index),
        )
        for index in range(len(LABELS))
    ]
    if (
        config.get("model_type") != "bert"
        or config.get("hidden_size") != 384
        or config.get("num_hidden_layers") != 6
    ):
        raise DataError("Core ML conversion requires the purpose-lite BERT architecture")
    if configured_labels != list(LABELS):
        raise DataError("Core ML checkpoint label order does not match purpose-lite")
    return config


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


def _package_manifest(package: Path) -> dict[str, Any]:
    files = []
    for path in sorted(item for item in package.rglob("*") if item.is_file()):
        files.append(
            {
                "path": str(path.relative_to(package)),
                "bytes": path.stat().st_size,
                "sha256": _sha256(path),
            }
        )
    return {
        "schemaVersion": 1,
        "modelVersion": "purpose-lite-v1",
        "package": package.name,
        "bytes": sum(item["bytes"] for item in files),
        "files": files,
    }


def _verify_wrapper_parity(
    torch_model: Any,
    wrapper: FixedShapeBertForCoreML,
    tokenizer: Any,
    records: Sequence[dict[str, Any]],
) -> float:
    encoded = encode_fixed_shape(
        tokenizer,
        [record["prompt"] for record in records],
        torch,
    )
    maximum_error = 0.0
    with torch.inference_mode():
        for index in range(len(records)):
            item = {key: value[index : index + 1] for key, value in encoded.items()}
            reference = torch_model(**item).logits
            candidate = wrapper(
                item["input_ids"].to(torch.int32),
                item["attention_mask"].to(torch.int32),
                item.get("token_type_ids", torch.zeros_like(item["input_ids"])).to(
                    torch.int32
                ),
            )
            error = float(torch.max(torch.abs(reference - candidate)).item())
            maximum_error = max(maximum_error, error)
            torch.testing.assert_close(candidate, reference, rtol=1e-5, atol=2e-5)
            if int(candidate.argmax(dim=-1).item()) != int(
                reference.argmax(dim=-1).item()
            ):
                raise DataError("conversion wrapper changed a parity prediction")
    return maximum_error


def convert(args: argparse.Namespace) -> dict[str, Any]:
    _checkpoint_config(args.model_dir)
    checkpoint = args.model_dir / "model.safetensors"
    if not checkpoint.is_file():
        raise DataError(f"{checkpoint}: checkpoint is missing")
    if args.output.suffix != ".mlpackage":
        raise DataError("Core ML output must end in .mlpackage")
    if args.output.exists():
        if not args.overwrite_output:
            raise DataError(
                f"{args.output}: output exists; pass --overwrite-output intentionally"
            )
        if args.output.is_dir():
            shutil.rmtree(args.output)
        else:
            args.output.unlink()
    args.output.parent.mkdir(parents=True, exist_ok=True)

    try:
        import coremltools as ct
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
    except ImportError as exc:
        raise DataError(
            "Core ML conversion requires requirements-coreml.txt on macOS"
        ) from exc
    if ct.__version__ != COREMLTOOLS_VERSION:
        raise DataError(
            f"expected coremltools {COREMLTOOLS_VERSION}, found {ct.__version__}"
        )

    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_dir,
        local_files_only=True,
    ).eval()
    tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True)
    wrapper = FixedShapeBertForCoreML(model).eval()
    validation = load_jsonl(args.validation)
    if len(validation) < args.parity_records:
        raise DataError("validation split is smaller than --parity-records")
    maximum_error = _verify_wrapper_parity(
        model,
        wrapper,
        tokenizer,
        validation[: args.parity_records],
    )
    print(f"conversion-wrapper parity: max_abs_error={maximum_error:.3g}")

    example = (
        torch.zeros((1, MAX_LENGTH), dtype=torch.int32),
        torch.ones((1, MAX_LENGTH), dtype=torch.int32),
        torch.zeros((1, MAX_LENGTH), dtype=torch.int32),
    )
    with torch.inference_mode():
        traced = torch.jit.trace(wrapper, example, strict=True)
        traced = torch.jit.freeze(traced)
    deployment_target = getattr(ct.target, args.minimum_deployment_target, None)
    if deployment_target is None:
        raise DataError(
            f"coremltools does not support {args.minimum_deployment_target}"
        )
    coreml_model = ct.convert(
        traced,
        convert_to="mlprogram",
        minimum_deployment_target=deployment_target,
        compute_precision=ct.precision.FLOAT16,
        inputs=[
            ct.TensorType(
                name="input_ids",
                shape=(1, MAX_LENGTH),
                dtype=np.int32,
            ),
            ct.TensorType(
                name="attention_mask",
                shape=(1, MAX_LENGTH),
                dtype=np.int32,
            ),
            ct.TensorType(
                name="token_type_ids",
                shape=(1, MAX_LENGTH),
                dtype=np.int32,
            ),
        ],
        outputs=[ct.TensorType(name="logits", dtype=np.float32)],
    )
    coreml_model.author = "Nucleic"
    coreml_model.short_description = "purpose-lite-v1 prompt classifier"
    coreml_model.version = "purpose-lite-v1"
    coreml_model.user_defined_metadata["com.nucleic.model.version"] = (
        "purpose-lite-v1"
    )
    coreml_model.user_defined_metadata["com.nucleic.model.labels"] = json.dumps(
        list(LABELS),
        separators=(",", ":"),
    )
    coreml_model.user_defined_metadata["com.nucleic.model.sourceSha256"] = _sha256(
        checkpoint
    )
    coreml_model.user_defined_metadata["com.nucleic.model.fixedShape"] = "1x128"
    coreml_model.user_defined_metadata["com.nucleic.model.minimumDeploymentTarget"] = (
        args.minimum_deployment_target
    )
    coreml_model.save(str(args.output))

    manifest = _package_manifest(args.output)
    manifest.update(
        {
            "sourceCheckpoint": str(checkpoint),
            "sourceCheckpointSha256": _sha256(checkpoint),
            "coremltoolsVersion": ct.__version__,
            "minimumDeploymentTarget": args.minimum_deployment_target,
            "computePrecision": "float16",
            "wrapperMaximumAbsoluteError": maximum_error,
        }
    )
    manifest_path = args.output.with_name(f"{args.output.stem}-manifest.json")
    write_json(manifest_path, manifest)
    print(f"Core ML package: {args.output} ({manifest['bytes']} bytes)")
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
    parser.add_argument("--validation", type=Path, default=DEFAULT_VALIDATION)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--parity-records", type=int, default=16)
    parser.add_argument(
        "--minimum-deployment-target",
        default="macOS15",
        choices=("macOS15", "macOS26"),
    )
    parser.add_argument("--overwrite-output", action="store_true")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.parity_records <= 0:
        parser.error("--parity-records must be positive")
    try:
        convert(args)
    except (AssertionError, DataError, OSError, RuntimeError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


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