from __future__ import annotations

import hashlib
import json
import math
import unicodedata
from pathlib import Path
from typing import Any, cast

import numpy as np
import skops.io as sio


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def normalize_text(value: str, maximum_characters: int) -> str:
    if not isinstance(value, str):
        raise TypeError("input must be a string")
    if len(value) > maximum_characters:
        raise ValueError(f"input must not exceed {maximum_characters} characters")
    normalized = " ".join(unicodedata.normalize("NFKC", value).casefold().split())
    if not normalized:
        raise ValueError("input must not be empty")
    return normalized


def risk_features(probabilities: np.ndarray, texts: list[str]) -> np.ndarray:
    clipped = np.clip(probabilities, 1e-12, 1.0)
    ordered = np.sort(clipped, axis=1)
    top1 = ordered[:, -1]
    top2 = ordered[:, -2]
    entropy = -(clipped * np.log(clipped)).sum(axis=1) / math.log(clipped.shape[1])
    shapes = np.asarray(
        [
            [
                min(len(text), 512) / 512,
                min(len(text.split()), 100) / 100,
                min(sum(ch.isdigit() for ch in text), 20) / 20,
                float("?" in text),
                float(
                    any(
                        token in text.split()
                        for token in ("not", "no", "never", "wrong")
                    )
                ),
            ]
            for text in texts
        ],
        dtype=np.float64,
    )
    predicted_one_hot = np.zeros_like(clipped)
    predicted_one_hot[np.arange(len(clipped)), np.argmax(clipped, axis=1)] = 1.0
    return np.column_stack(
        [clipped, predicted_one_hot, top1, top2, top1 - top2, entropy, shapes]
    )


def _load_checked(path: Path, expected_hash: str) -> dict[str, Any]:
    if sha256_file(path) != expected_hash:
        raise RuntimeError(f"artifact checksum mismatch: {path.name}")
    untrusted = sio.get_untrusted_types(file=path)
    if untrusted:
        raise RuntimeError(f"artifact requires untrusted skops types: {untrusted}")
    return cast(dict[str, Any], sio.load(path, trusted=[]))


def predict_text(
    text: str,
    *,
    primary_path: Path = Path("primary_baseline.skops"),
    risk_path: Path = Path("model.skops"),
    config_path: Path = Path("config.json"),
) -> dict[str, Any]:
    config = cast(dict[str, Any], json.loads(config_path.read_text(encoding="utf-8")))
    if config.get("schema_version") != "1.0":
        raise RuntimeError("unsupported configuration schema")
    normalized = normalize_text(text, int(config["input"]["maximum_characters"]))
    primary = _load_checked(primary_path, str(config["primary_artifact_sha256"]))
    candidate = _load_checked(risk_path, str(config["artifact_sha256"]))
    labels = cast(list[str], config["labels"])
    if labels != primary["labels"] or labels != candidate["labels"]:
        raise RuntimeError("label mapping mismatch")
    probabilities = np.asarray(
        primary["classifier"].predict_proba(
            primary["features"].transform([normalized])
        ),
        dtype=np.float64,
    )
    if probabilities.shape != (1, len(labels)) or not np.isfinite(probabilities).all():
        raise RuntimeError("primary model returned invalid probabilities")
    predicted = int(probabilities[0].argmax())
    error_risk = float(
        candidate["classifier"].predict_proba(
            risk_features(probabilities, [normalized])
        )[0, 1]
    )
    threshold = float(config["review_at_or_above_error_risk"])
    return {
        "intent": labels[predicted],
        "intent_score": float(probabilities[0, predicted]),
        "error_risk": error_risk,
        "review": error_risk >= threshold,
        "review_threshold": threshold,
        "advisory_only": True,
        "score_semantics": config["score_semantics"],
    }


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="Route one banking-support query")
    parser.add_argument("text")
    parser.add_argument("--primary", type=Path, default=Path("primary_baseline.skops"))
    parser.add_argument("--risk-model", type=Path, default=Path("model.skops"))
    parser.add_argument("--config", type=Path, default=Path("config.json"))
    args = parser.parse_args()
    result = predict_text(
        args.text,
        primary_path=args.primary,
        risk_path=args.risk_model,
        config_path=args.config,
    )
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()