from __future__ import annotations

import json
from pathlib import Path
from typing import Literal, Sequence, TypedDict, cast

import numpy as np
import onnxruntime as ort
from PIL import Image, ImageOps


MODEL_DIR = Path(__file__).resolve().parent.parent
PADDING_MULTIPLE = 32
ModelFormat = Literal["fp32", "fp16"]
MODEL_FILENAMES: dict[ModelFormat, str] = {
    "fp32": "model.onnx",
    "fp16": "model.fp16.onnx",
}


class Preprocess(TypedDict):
    resize_longest_side_px: int
    mean: list[float]
    std: list[float]


class LabelSet(TypedDict):
    labels: list[str]


class Labels(TypedDict):
    screen: LabelSet
    safety: LabelSet


class Prediction(TypedDict):
    screen: str
    safety: str


class Classifier:
    def __init__(
        self,
        model_dir: str | Path = MODEL_DIR,
        providers: Sequence[str] = ("CPUExecutionProvider",),
        model_format: ModelFormat = "fp32",
    ) -> None:
        self.directory = Path(model_dir)
        self.session = ort.InferenceSession(str(model_path(self.directory, model_format)), providers=list(providers))
        self.preprocess = load_preprocess(self.directory / "preprocess.json")
        self.labels = load_labels(self.directory / "inference" / "labels.json")

    def classify(self, image_path: str | Path) -> Prediction:
        return self.classify_batch([image_path])[0]

    def classify_batch(self, image_paths: Sequence[str | Path]) -> list[Prediction]:
        if not image_paths:
            return []
        images = [preprocess_image(Path(image_path), self.preprocess) for image_path in image_paths]
        screen_logits, safety_logits = self.session.run(None, {"image": collate_images(images)})
        return decode_predictions(screen_logits, safety_logits, self.labels)


def classify(
    image_path: str | Path,
    model_dir: str | Path = MODEL_DIR,
    model_format: ModelFormat = "fp32",
) -> Prediction:
    return Classifier(model_dir, model_format=model_format).classify(image_path)


def classify_batch(
    image_paths: Sequence[str | Path],
    model_dir: str | Path = MODEL_DIR,
    model_format: ModelFormat = "fp32",
) -> list[Prediction]:
    return Classifier(model_dir, model_format=model_format).classify_batch(image_paths)


def model_path(directory: Path, model_format: ModelFormat) -> Path:
    path = directory / "onnx" / MODEL_FILENAMES[model_format]
    if not path.is_file():
        raise FileNotFoundError(f"Missing {model_format} ONNX model: {path}")
    return path


def preprocess_image(image_path: Path, preprocess: Preprocess) -> np.ndarray:
    with Image.open(image_path) as opened:
        image = to_training_rgb(opened)
    resized = resize_image(image, preprocess["resize_longest_side_px"])
    array = np.asarray(resized).astype("float32") / 255.0
    mean = np.asarray(preprocess["mean"], dtype="float32")
    std = np.asarray(preprocess["std"], dtype="float32")
    array = (array - mean) / std
    return np.transpose(array, (2, 0, 1))


def to_training_rgb(image: Image.Image) -> Image.Image:
    image = ImageOps.exif_transpose(image)
    if image.mode == "P" and isinstance(image.info.get("transparency"), bytes):
        image = image.convert("RGBA")
    if image.mode in ("RGBA", "LA", "PA"):
        rgba = image.convert("RGBA")
        background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
        image = Image.alpha_composite(background, rgba)
    return image.convert("RGB")


def collate_images(images: Sequence[np.ndarray]) -> np.ndarray:
    height = round_up(max(int(image.shape[1]) for image in images))
    width = round_up(max(int(image.shape[2]) for image in images))
    batch = np.zeros((len(images), 3, height, width), dtype="float32")
    for index, image in enumerate(images):
        image_height = int(image.shape[1])
        image_width = int(image.shape[2])
        batch[index, :, :image_height, :image_width] = image
    return batch


def resize_image(image: Image.Image, image_size: int) -> Image.Image:
    scale = image_size / max(image.width, image.height)
    width = max(1, round(image.width * scale))
    height = max(1, round(image.height * scale))
    return image.resize((width, height), Image.Resampling.BICUBIC)


def round_up(value: int, multiple: int = PADDING_MULTIPLE) -> int:
    return ((value + multiple - 1) // multiple) * multiple


def decode_predictions(screen_logits: np.ndarray, safety_logits: np.ndarray, labels: Labels) -> list[Prediction]:
    # ONNX emits flat screen and safety logits.
    screen_indices = top_indices(screen_logits)
    safety_indices = top_indices(safety_logits)
    return [
        {
            "screen": labels["screen"]["labels"][screen_index],
            "safety": labels["safety"]["labels"][safety_index],
        }
        for screen_index, safety_index in zip(screen_indices, safety_indices, strict=True)
    ]


def load_labels(path: Path) -> Labels:
    return cast(Labels, json.loads(path.read_text(encoding="utf-8")))


def load_preprocess(path: Path) -> Preprocess:
    return cast(Preprocess, json.loads(path.read_text(encoding="utf-8")))


def top_indices(logits: np.ndarray) -> list[int]:
    return [int(index) for index in np.argmax(logits, axis=1)]


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Classify images with the exported ONNX model.")
    parser.add_argument("images", nargs="+")
    parser.add_argument("--model-dir", default=str(MODEL_DIR))
    parser.add_argument("--model-format", choices=tuple(MODEL_FILENAMES), default="fp32")
    args = parser.parse_args()
    predictions = classify_batch(args.images, args.model_dir, cast(ModelFormat, args.model_format))
    value: Prediction | list[Prediction] = predictions[0] if len(predictions) == 1 else predictions
    print(json.dumps(value, indent=2))