"""inpaintVideo: LaMa (Apache-2.0, ONNX) per frame with optical-flow smoothing.

Per frame: crop around the mask, inpaint at 512x512, paste back through a
feathered mask. When the previous frame was inpainted too, its result is
warped into the current frame with DIS optical flow and blended in, which
damps flicker. Quality is measured on the output:
- temporalStability: 1 - mean absolute change inside the mask between
  consecutive frames after motion compensation (scaled so 32 levels -> 0).
- edgeQuality: 1 - excess gradient on the mask boundary over the gradient
  just outside it.
- score: 0.6 * temporalStability + 0.4 * edgeQuality.
Callers fall back to cover/reframe when the score is low.
"""

from __future__ import annotations

import os
import threading
from typing import Any

import cv2
import numpy as np

from .. import media_io, paths
from ..errors import BadRequest
from ..schemas import InpaintVideoRequest
from .common import Ctx, clamp01, open_input
from .devices import onnx_providers
from .masks import MaskIndex
from .redact import OUTPUT_EXTENSIONS

LAMA_SIZE = 512
TEMPORAL_BLEND = 0.35


class LamaInpainter:
    def __init__(self, path: str, settings: Any) -> None:
        import onnxruntime

        self.session = onnxruntime.InferenceSession(path, providers=onnx_providers(settings))

    def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """image: BGR uint8 512x512, mask: uint8 0/255 512x512 -> BGR uint8."""
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        pixel = rgb.transpose(2, 0, 1)[None]
        m = (mask > 127).astype(np.float32)[None, None]
        (out,) = self.session.run(["output"], {"image": pixel, "mask": m})
        result = np.clip(out[0].transpose(1, 2, 0), 0, 255).astype(np.uint8)
        return cv2.cvtColor(result, cv2.COLOR_RGB2BGR)


def load_inpainter(files: dict[str, str], settings: Any) -> LamaInpainter:
    return LamaInpainter(files["lama_fp32.onnx"], settings)


def _crop_box(mask: np.ndarray) -> tuple[int, int, int, int]:
    height, width = mask.shape
    ys, xs = np.nonzero(mask)
    cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2
    side = max(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1) * 2.2
    side = int(min(max(side, 256), max(width, height)))
    x0 = int(min(max(cx - side / 2, 0), max(width - side, 0)))
    y0 = int(min(max(cy - side / 2, 0), max(height - side, 0)))
    return x0, y0, min(x0 + side, width), min(y0 + side, height)


class Smoother:
    """Temporal blending and quality bookkeeping across frames."""

    def __init__(self) -> None:
        self.flow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_FAST)
        self.prev_gray: np.ndarray | None = None
        self.prev_out: np.ndarray | None = None
        self.prev_masked = False
        self.stability: list[float] = []
        self.edges: list[float] = []

    def _warp_prev(self, gray: np.ndarray) -> np.ndarray | None:
        if self.prev_gray is None or self.prev_out is None:
            return None
        flow = self.flow.calc(gray, self.prev_gray, None)
        h, w = gray.shape
        grid_x, grid_y = np.meshgrid(np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32))
        return cv2.remap(
            self.prev_out,
            grid_x + flow[..., 0],
            grid_y + flow[..., 1],
            cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_REPLICATE,
        )

    def step(self, source: np.ndarray, inpainted: np.ndarray | None, mask: np.ndarray | None) -> np.ndarray:
        gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY)
        out = source
        if inpainted is not None and mask is not None:
            soft = cv2.GaussianBlur(mask.astype(np.float32) / 255.0, (0, 0), 3)
            soft = np.maximum(soft, mask.astype(np.float32) / 255.0)[..., None]
            warped = self._warp_prev(gray) if self.prev_masked else None
            fill = inpainted.astype(np.float32)
            if warped is not None:
                fill = fill * (1 - TEMPORAL_BLEND) + warped.astype(np.float32) * TEMPORAL_BLEND
            out = np.clip(source.astype(np.float32) * (1 - soft) + fill * soft, 0, 255).astype(np.uint8)
            inside = mask > 0
            if warped is not None:
                # Raw LaMa output against the motion-compensated previous frame: flicker before smoothing.
                diff = np.abs(inpainted.astype(np.float32) - warped.astype(np.float32))[inside].mean()
                self.stability.append(clamp01(1.0 - float(diff) / 32.0))
            self.edges.append(_edge_quality(out, mask))
        self.prev_gray = gray
        self.prev_out = out
        self.prev_masked = mask is not None
        return out

    def quality(self) -> dict[str, float]:
        temporal = float(np.mean(self.stability)) if self.stability else 1.0
        edge = float(np.mean(self.edges)) if self.edges else 1.0
        return {
            "temporalStability": clamp01(temporal),
            "edgeQuality": clamp01(edge),
            "score": clamp01(0.6 * temporal + 0.4 * edge),
        }


def _edge_quality(frame: np.ndarray, mask: np.ndarray) -> float:
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32)
    grad = cv2.magnitude(cv2.Sobel(gray, cv2.CV_32F, 1, 0), cv2.Sobel(gray, cv2.CV_32F, 0, 1))
    kernel = np.ones((5, 5), np.uint8)
    boundary = (cv2.dilate(mask, kernel) > 0) & (cv2.erode(mask, kernel) == 0)
    outer = (cv2.dilate(mask, np.ones((15, 15), np.uint8)) > 0) & ~(cv2.dilate(mask, kernel) > 0)
    if not boundary.any() or not outer.any():
        return 1.0
    excess = max(0.0, float(grad[boundary].mean()) - float(grad[outer].mean()))
    return clamp01(1.0 - excess / (float(grad[outer].mean()) + 40.0))


def run(ctx: Ctx, body: InpaintVideoRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    masks = MaskIndex(paths.input_dir(ctx.settings.media_root, body.maskDir), info.frame_ms())
    out_real = paths.output_file(ctx.settings.media_root, body.outPath)
    if os.path.splitext(out_real)[1].lower() not in OUTPUT_EXTENSIONS:
        raise BadRequest("outPath must end in .mkv, .mp4 or .mov")
    if out_real == real:
        raise BadRequest("outPath must differ from path")
    inpainter = ctx.registry.runtime("inpaintVideo")
    smoother = Smoother()

    def per_frame(_index: int, t_ms: int, frame: np.ndarray) -> np.ndarray:
        height, width = frame.shape[:2]
        mask = masks.lookup(t_ms, width, height)
        if mask is None or not mask.any():
            return smoother.step(frame, None, None)
        # Grow slightly so the seam lands on untouched pixels.
        mask = cv2.dilate(mask, np.ones((7, 7), np.uint8))
        x0, y0, x1, y1 = _crop_box(mask)
        crop = cv2.resize(frame[y0:y1, x0:x1], (LAMA_SIZE, LAMA_SIZE), interpolation=cv2.INTER_AREA)
        crop_mask = cv2.resize(mask[y0:y1, x0:x1], (LAMA_SIZE, LAMA_SIZE), interpolation=cv2.INTER_NEAREST)
        filled = inpainter.inpaint(crop, crop_mask)
        full = frame.copy()
        full[y0:y1, x0:x1] = cv2.resize(filled, (x1 - x0, y1 - y0), interpolation=cv2.INTER_CUBIC)
        return smoother.step(frame, full, mask)

    media_io.transform_video(real, out_real, info, cancel, per_frame)
    return {"outPath": body.outPath, "quality": smoother.quality()}