"""Serve unlabelled Kev 0.5B typed decisions from the Core ML artifact."""

from __future__ import annotations

import argparse
import json
import shutil
import tempfile
from pathlib import Path

import coremltools as ct
import numpy as np
from kev.api import output_tokens, to_answers
from transformers import AutoTokenizer

from preprocessing import Shape, prepare_runtime_inputs

COMPUTE_UNITS = {
    "all": ct.ComputeUnit.ALL,
    "cpu": ct.ComputeUnit.CPU_ONLY,
    "cpu-gpu": ct.ComputeUnit.CPU_AND_GPU,
    "cpu-ne": ct.ComputeUnit.CPU_AND_NE,
}
PACKAGE = "kev_0_5b_fp16_L128_options32.mlpackage"


def stage_symlinked_package(package: Path) -> tuple[Path, tempfile.TemporaryDirectory | None]:
    """Give Core ML real package files when a Hub snapshot contains blob symlinks."""
    if not package.is_symlink() and not any(path.is_symlink() for path in package.rglob("*")):
        return package, None
    temporary = tempfile.TemporaryDirectory(prefix="kev-coreml-")
    staged = Path(temporary.name) / package.name
    shutil.copytree(package, staged, symlinks=False)
    return staged, temporary


class KevCoreML:
    """Reusable serving session; initialization and request latency stay separate."""

    def __init__(self, model_dir: str | Path, *, units: str = "all", package: str | Path | None = None):
        root = Path(model_dir)
        self.tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer", local_files_only=True)
        self.shape = Shape()
        model_path = Path(package) if package is not None else root / PACKAGE
        model_path, self._staged_package = stage_symlinked_package(model_path)
        self.model = ct.models.MLModel(str(model_path), compute_units=COMPUTE_UNITS[units])

    def predict(self, request: dict) -> dict:
        arrays, encoded, metadata, parsed = prepare_runtime_inputs(self.tokenizer, request, self.shape)
        result = np.asarray(self.model.predict(arrays)["probabilities"], dtype=np.float64)
        count = len(metadata[0]["keys"])
        if result.shape != (1, self.shape.max_options) or not np.isfinite(result).all():
            raise ValueError("Invalid Core ML probabilities")
        selected = result[0, :count].tolist()
        keys = metadata[0]["keys"]
        answers = to_answers([selected], metadata)
        return {
            "model": parsed.model,
            "answers": answers,
            "usage": {
                "input_tokens": len(encoded["ids"]),
                "output_tokens": output_tokens(self.tokenizer, answers),
            },
            "option_keys": keys,
            "probabilities": selected,
            "chosen": keys[int(np.argmax(selected))],
        }


def predict(model_dir: str | Path, request: dict, *, units: str = "all", package: str | Path | None = None) -> dict:
    """One-shot compatibility wrapper; reuse `KevCoreML` for repeated calls."""
    return KevCoreML(model_dir, units=units, package=package).predict(request)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-dir", required=True, type=Path)
    parser.add_argument("--request-json", type=Path, required=True)
    parser.add_argument("--units", choices=COMPUTE_UNITS, default="all")
    parser.add_argument("--package", type=Path)
    args = parser.parse_args()
    request = json.loads(args.request_json.read_text())
    print(json.dumps(predict(args.model_dir, request, units=args.units, package=args.package), indent=2))


if __name__ == "__main__":
    main()