"""WD14 EVA02-Large tagger (ONNX CPU inference)."""
from __future__ import annotations
import csv
import os
from typing import Iterable

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

# Some images in the library are minutely truncated (e.g. 6 bytes short of
# the JPEG end-of-image marker). The model doesn't care about the last few
# pixels, but PIL's default strict load raises OSError. Tolerate them so a
# single corrupt tail doesn't block tagging the rest of the image.
ImageFile.LOAD_TRUNCATED_IMAGES = True

MODEL_VERSION = os.environ.get('WD14_MODEL_VERSION', 'wd-eva02-large-tagger-v3')
_MODEL_DIR = os.environ.get('ML_MODEL_DIR', '/models')
_WD14_DIR = os.path.join(_MODEL_DIR, 'wd14')
_MODEL_PATH = os.path.join(_WD14_DIR, 'model.onnx')
_TAGS_PATH = os.path.join(_WD14_DIR, 'selected_tags.csv')

# WD14 selected_tags.csv uses Danbooru category ids:
#   0=general, 1=artist, 3=copyright, 4=character, 5=meta, 9=rating
_CATEGORY_MAP = {0: 'general', 1: 'artist', 3: 'copyright', 4: 'character', 5: 'meta', 9: 'rating'}

_session: ort.InferenceSession | None = None
_tag_meta: list[dict] | None = None
_input_name: str | None = None
_output_name: str | None = None
_input_size: int = 448


# NOT thread-safe. Must run in the ml-worker container with --concurrency=1.
def _load() -> None:
    global _session, _tag_meta, _input_name, _output_name, _input_size
    if _session is not None:
        return

    if not os.path.isfile(_MODEL_PATH):
        raise RuntimeError(
            f"WD14 model file missing at {_MODEL_PATH}. "
            f"Populate the /models volume via the ml-worker downloader."
        )
    if not os.path.isfile(_TAGS_PATH):
        raise RuntimeError(
            f"WD14 selected_tags.csv missing at {_TAGS_PATH}. "
            f"Populate the /models volume via the ml-worker downloader."
        )

    tag_meta: list[dict] = []
    with open(_TAGS_PATH, newline='') as f:
        reader = csv.DictReader(f)
        for row in reader:
            tag_meta.append({
                'name': row['name'],
                'category': _CATEGORY_MAP.get(int(row['category']), 'unknown'),
            })

    session = ort.InferenceSession(
        _MODEL_PATH,
        providers=['CPUExecutionProvider'],
    )
    _input_name = session.get_inputs()[0].name
    _output_name = session.get_outputs()[0].name
    # Input shape is usually [batch, H, W, 3] NHWC; pick the spatial dim
    input_shape = session.get_inputs()[0].shape
    for dim in input_shape:
        if isinstance(dim, int) and dim > 1:
            _input_size = dim
            break
    # Assign sentinels last so a partially-loaded state can't be observed.
    _tag_meta = tag_meta
    _session = session


def _preprocess(image_path: str) -> np.ndarray:
    img = Image.open(image_path)
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    # Composite onto white background so transparency doesn't bias the model
    bg = Image.new('RGBA', img.size, (255, 255, 255, 255))
    bg.paste(img, mask=img.split()[3] if img.mode == 'RGBA' else None)
    img = bg.convert('RGB')

    w, h = img.size
    side = max(w, h)
    square = Image.new('RGB', (side, side), (255, 255, 255))
    square.paste(img, ((side - w) // 2, (side - h) // 2))
    square = square.resize((_input_size, _input_size), Image.BICUBIC)

    arr = np.array(square, dtype=np.float32)
    # WD14 was trained on BGR
    arr = arr[:, :, ::-1]
    return arr[np.newaxis, :, :, :]  # NHWC


def infer(image_path: str) -> list[dict]:
    """Run WD14 on one image. Returns a list of {name, category, confidence}."""
    _load()
    x = _preprocess(image_path)
    out = _session.run([_output_name], {_input_name: x})[0][0]
    results: list[dict] = []
    for idx, score in enumerate(out):
        meta = _tag_meta[idx]
        results.append({
            'name': meta['name'],
            'category': meta['category'],
            'confidence': float(score),
        })
    return results


def infer_filtered(image_path: str, min_any: float = 0.05) -> list[dict]:
    """Same as infer() but drops tags below a floor to keep DB rows reasonable."""
    return [r for r in infer(image_path) if r['confidence'] >= min_any]