# -*- coding: utf-8 -*-
import json
import logging
import os
from typing import List, Optional, Tuple

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

logger = logging.getLogger(__name__)

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


class OnnxClassifier:
    def __init__(self, model_path: str, image_size: int = 224):
        self.image_size = image_size
        self.session: Optional[ort.InferenceSession] = None
        if model_path and os.path.isfile(model_path):
            self.session = ort.InferenceSession(
                model_path,
                providers=["CPUExecutionProvider"],
            )
            logger.info("加载 ONNX: %s", model_path)

    @property
    def loaded(self) -> bool:
        return self.session is not None

    def preprocess(self, image_path: str) -> np.ndarray:
        img = Image.open(image_path).convert("RGB").resize((self.image_size, self.image_size))
        arr = np.array(img).astype(np.float32) / 255.0
        arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
        return arr.transpose(2, 0, 1)[None]

    def predict_batch(self, image_paths: List[str], batch_size: int = 16) -> List[float]:
        if not self.session:
            return [0.0] * len(image_paths)
        scores = []
        for i in range(0, len(image_paths), batch_size):
            batch_paths = image_paths[i : i + batch_size]
            batch = np.concatenate([self.preprocess(p) for p in batch_paths], axis=0)
            logits = self.session.run(None, {"input": batch})[0]
            for logit in logits:
                scores.append(float(1.0 / (1.0 + np.exp(-logit[0]))))
        return scores


class ModelRegistry:
    def __init__(self, model_dir: str):
        self.model_dir = model_dir
        self.version = "unknown"
        self.image_size = 224
        self.storefront: Optional[OnnxClassifier] = None
        self.handover: Optional[OnnxClassifier] = None
        self._load()

    def _load(self):
        version_path = os.path.join(self.model_dir, "version.json")
        sf_name = ho_name = None
        if os.path.isfile(version_path):
            with open(version_path, encoding="utf-8") as f:
                meta = json.load(f)
            self.version = meta.get("model_version", "1.0.0")
            self.image_size = int(meta.get("image_size", 224))
            sf_name = meta.get("storefront_model")
            ho_name = meta.get("handover_model")

        def resolve(name: str, fallback: str) -> str:
            if name:
                path = os.path.join(self.model_dir, name)
                if os.path.isfile(path):
                    return path
            for suffix in (f"{fallback}_int8.onnx", f"{fallback}.onnx"):
                path = os.path.join(self.model_dir, suffix)
                if os.path.isfile(path):
                    return path
            return ""

        sf_path = resolve(sf_name, "storefront")
        ho_path = resolve(ho_name, "handover")
        self.storefront = OnnxClassifier(sf_path, self.image_size) if sf_path else OnnxClassifier("", self.image_size)
        self.handover = OnnxClassifier(ho_path, self.image_size) if ho_path else OnnxClassifier("", self.image_size)

    @property
    def ready(self) -> bool:
        return self.storefront.loaded and self.handover.loaded


def score_frames(
    registry: ModelRegistry,
    frames: List[Tuple[float, str]],
) -> Tuple[List[Tuple[float, float]], List[Tuple[float, float]]]:
    times = [f[0] for f in frames]
    paths = [f[1] for f in frames]
    sf_scores = registry.storefront.predict_batch(paths)
    ho_scores = registry.handover.predict_batch(paths)
    return list(zip(times, sf_scores)), list(zip(times, ho_scores))