"""ProPainter wrapper with MPS adaptation, following official inference pipeline."""

from __future__ import annotations

import sys
import time
from pathlib import Path
from typing import Any, Callable

import cv2
import numpy as np
from loguru import logger


class _nullcontext:
    """Minimal no-op context manager (avoids importing contextlib)."""
    def __enter__(self): return self
    def __exit__(self, *a): pass


class InpaintTimeoutError(RuntimeError):
    pass


class ProPainterEngine:
    """Wraps ProPainter for MPS-compatible video inpainting.

    Follows the official inference_propainter.py pipeline:
    1. RAFT optical flow
    2. FlowComplete bidirectional flow completion
    3. Image propagation
    4. InpaintGenerator (sliding window transformer)
    """

    TIMEOUT_1080P = 30
    TIMEOUT_4K = 60
    MAX_CONSECUTIVE_TIMEOUTS = 5

    def __init__(
        self,
        device: str = "mps",
        neighbor_length: int = 10,
        ref_stride: int = 10,
        propainter_root: str | None = None,
        use_fp16: bool = True,
    ) -> None:
        self.device = device
        self.neighbor_length = neighbor_length
        self.ref_stride = ref_stride
        self.use_fp16 = use_fp16
        self.propainter_root = propainter_root or str(
            Path(__file__).parent.parent / "third_party" / "ProPainter"
        )

        self._raft_model: Any = None
        self._flow_comp_model: Any = None
        self._inpaint_model: Any = None
        self._loaded = False
        self._patched = False

    def load_models(self) -> None:
        if self._loaded:
            return

        import torch

        pp_path = Path(self.propainter_root)
        if not pp_path.exists():
            raise FileNotFoundError(f"ProPainter not found at {pp_path}")

        if str(pp_path) not in sys.path:
            sys.path.insert(0, str(pp_path))

        self._patch_mps_compat()

        from model.modules.flow_comp_raft import RAFT_bi
        from model.recurrent_flow_completion import RecurrentFlowCompleteNet
        from model.propainter import InpaintGenerator

        weights_dir = pp_path / "weights"

        self._raft_model = RAFT_bi(str(weights_dir / "raft-things.pth"), device=self.device)

        self._flow_comp_model = RecurrentFlowCompleteNet(str(weights_dir / "recurrent_flow_completion.pth"))
        for p in self._flow_comp_model.parameters():
            p.requires_grad = False
        self._flow_comp_model.to(self.device)

        self._inpaint_model = InpaintGenerator(model_path=str(weights_dir / "ProPainter.pth")).to(self.device)
        self._inpaint_model.requires_grad_(False)

        self._loaded = True
        logger.info(f"ProPainter models loaded on {self.device}")

    def _patch_mps_compat(self) -> None:
        if self._patched:
            return

        import torch

        try:
            import torchvision.ops as tv_ops
            original_deform_conv = tv_ops.deform_conv2d

            def mps_safe_deform_conv(*args, **kwargs):
                cpu_args = [a.cpu() if isinstance(a, torch.Tensor) else a for a in args]
                cpu_kwargs = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
                result = original_deform_conv(*cpu_args, **cpu_kwargs)
                return result.to("mps")

            tv_ops.deform_conv2d = mps_safe_deform_conv
            logger.debug("Patched deform_conv2d for MPS CPU fallback")
        except Exception as e:
            logger.warning(f"Failed to patch deform_conv2d: {e}")

        if hasattr(torch.cuda, "synchronize"):
            torch.cuda.synchronize = lambda device=None: torch.mps.synchronize()

        # Patch torch.cuda.empty_cache for ProPainter code that calls it directly
        if hasattr(torch.cuda, "empty_cache"):
            original_empty = torch.cuda.empty_cache
            def mps_empty_cache():
                torch.mps.empty_cache()
            torch.cuda.empty_cache = mps_empty_cache

        self._patched = True

    def inpaint(
        self,
        frames: list[np.ndarray],
        masks: list[np.ndarray],
        frame_timeout: float | None = None,
    ) -> list[np.ndarray]:
        """Run ProPainter inpainting following official pipeline.

        Args:
            frames: list of (H, W, 3) uint8 RGB
            masks: list of (H, W) bool, same length as frames

        Returns:
            list of inpainted (H, W, 3) uint8 RGB frames
        """
        if not self._loaded:
            raise RuntimeError("Models not loaded. Call load_models() first.")

        import torch
        import torch.nn.functional as F

        n = len(frames)
        if n == 0:
            return []

        h, w = frames[0].shape[:2]

        # Convert frames to tensor: (1, T, 3, H, W), range [-1, 1]
        frames_t = torch.stack([
            torch.from_numpy(f).permute(2, 0, 1).float() / 255.0
            for f in frames
        ]).unsqueeze(0).to(self.device)
        frames_t = frames_t * 2 - 1  # normalize to [-1, 1]

        # Masks: (1, T, 1, H, W), values 0 or 1
        masks_t = torch.stack([
            torch.from_numpy(m.astype(np.float32)).unsqueeze(0)
            for m in masks
        ]).unsqueeze(0).to(self.device)

        # Flow masks = same as masks for inpainting mode
        flow_masks = masks_t.clone()

        # Use float16 autocast on MPS for ~1.5-2x speedup + lower memory
        autocast_ctx = (
            torch.autocast(self.device, dtype=torch.float16)
            if self.use_fp16 and self.device == "mps"
            else _nullcontext()
        )

        with torch.no_grad(), autocast_ctx:
            # --- Step 1: RAFT optical flow ---
            logger.debug("Computing optical flow...")
            short_clip_len = 12 if w <= 640 else (8 if w <= 720 else (4 if w <= 1280 else 2))

            if n > short_clip_len:
                gt_flows_f_list, gt_flows_b_list = [], []
                for f_idx in range(0, n, short_clip_len):
                    end_f = min(n, f_idx + short_clip_len)
                    if f_idx == 0:
                        fl_f, fl_b = self._raft_model(frames_t[:, f_idx:end_f], iters=20)
                    else:
                        fl_f, fl_b = self._raft_model(frames_t[:, f_idx-1:end_f], iters=20)
                        fl_f = fl_f[:, 1:]
                        fl_b = fl_b[:, 1:]
                    gt_flows_f_list.append(fl_f)
                    gt_flows_b_list.append(fl_b)
                    torch.mps.empty_cache()
                gt_flows_f = torch.cat(gt_flows_f_list, dim=1)
                gt_flows_b = torch.cat(gt_flows_b_list, dim=1)
            else:
                gt_flows_f, gt_flows_b = self._raft_model(frames_t, iters=20)

            # Ensure exactly n-1 flows (pad or trim)
            target_len = n - 1
            if gt_flows_f.shape[1] < target_len:
                pad = target_len - gt_flows_f.shape[1]
                gt_flows_f = torch.cat([gt_flows_f, gt_flows_f[:, -1:].repeat(1, pad, 1, 1, 1)], dim=1)
                gt_flows_b = torch.cat([gt_flows_b, gt_flows_b[:, -1:].repeat(1, pad, 1, 1, 1)], dim=1)
            gt_flows_f = gt_flows_f[:, :target_len]
            gt_flows_b = gt_flows_b[:, :target_len]

            gt_flows_bi = (gt_flows_f, gt_flows_b)

            # --- Step 2: Flow completion ---
            # FlowComplete expects flow_masks with T frames (not T-1)
            # It internally handles the dimension alignment
            logger.debug("Completing flow...")
            pred_flows_bi, _ = self._flow_comp_model.forward_bidirect_flow(gt_flows_bi, masks_t)
            pred_flows_bi = self._flow_comp_model.combine_flow(gt_flows_bi, pred_flows_bi, masks_t)
            torch.mps.empty_cache()

            # --- Step 3: Image propagation ---
            logger.debug("Image propagation...")
            masked_frames = frames_t * (1 - masks_t)
            prop_imgs, updated_local_masks = self._inpaint_model.img_propagation(
                masked_frames, pred_flows_bi, masks_t, 'nearest'
            )
            b, t = 1, n
            updated_frames = frames_t * (1 - masks_t) + prop_imgs.view(b, t, 3, h, w) * masks_t
            updated_masks = updated_local_masks.view(b, t, 1, h, w)
            torch.mps.empty_cache()

            # --- Step 4: InpaintGenerator (sliding window) ---
            logger.debug("Inpainting with transformer...")
            ori_frames_np = [np.array(f) for f in frames]
            comp_frames = [None] * n

            neighbor_stride = self.neighbor_length // 2
            ref_num = -1 if n <= 80 else (80 // self.ref_stride)

            for f_idx in range(0, n, neighbor_stride):
                neighbor_ids = list(range(
                    max(0, f_idx - neighbor_stride),
                    min(n, f_idx + neighbor_stride + 1),
                ))
                ref_ids = self._get_ref_index(f_idx, neighbor_ids, n, self.ref_stride, ref_num)

                selected_imgs = updated_frames[:, neighbor_ids + ref_ids]
                selected_masks = masks_t[:, neighbor_ids + ref_ids]
                selected_update_masks = updated_masks[:, neighbor_ids + ref_ids]
                selected_flows_bi = (
                    pred_flows_bi[0][:, [i for i in neighbor_ids[:-1]]],
                    pred_flows_bi[1][:, [i for i in neighbor_ids[:-1]]],
                )

                l_t = len(neighbor_ids)
                pred_img = self._inpaint_model(
                    selected_imgs, selected_flows_bi,
                    selected_masks, selected_update_masks, l_t,
                )
                pred_img = pred_img.view(-1, 3, h, w)
                pred_img = (pred_img + 1) / 2
                pred_img = pred_img.cpu().permute(0, 2, 3, 1).numpy() * 255

                binary_masks = masks_t[0, neighbor_ids].cpu().permute(0, 2, 3, 1).numpy().astype(np.uint8)

                for i, idx in enumerate(neighbor_ids):
                    img = pred_img[i].astype(np.uint8) * binary_masks[i] + \
                          ori_frames_np[idx] * (1 - binary_masks[i])
                    if comp_frames[idx] is None:
                        comp_frames[idx] = img
                    else:
                        comp_frames[idx] = (comp_frames[idx].astype(np.float32) * 0.5 +
                                            img.astype(np.float32) * 0.5)
                    comp_frames[idx] = comp_frames[idx].astype(np.uint8)

                torch.mps.empty_cache()

        # Fill any None frames with original
        result = []
        for i in range(n):
            if comp_frames[i] is not None:
                result.append(comp_frames[i])
            else:
                result.append(frames[i])

        return result

    @staticmethod
    def _get_ref_index(mid: int, neighbor_ids: list[int], length: int, ref_stride: int, ref_num: int) -> list[int]:
        """Get reference frame indices for transformer attention."""
        ref_index = []
        if ref_num == -1:
            for i in range(0, length, ref_stride):
                if i not in neighbor_ids:
                    ref_index.append(i)
        else:
            start_idx = max(0, mid - ref_stride * (ref_num // 2))
            end_idx = min(length, mid + ref_stride * (ref_num // 2))
            for i in range(start_idx, end_idx, ref_stride):
                if i not in neighbor_ids:
                    if len(ref_index) >= ref_num:
                        break
                    ref_index.append(i)
        return ref_index

    def _to_cpu(self, obj: Any) -> Any:
        try:
            import torch
            if isinstance(obj, torch.Tensor):
                return obj.cpu()
        except ImportError:
            pass
        if isinstance(obj, (tuple, list)):
            return type(obj)(self._to_cpu(x) for x in obj)
        if isinstance(obj, dict):
            return {k: self._to_cpu(v) for k, v in obj.items()}
        return obj

    def _to_device(self, obj: Any) -> Any:
        try:
            import torch
            if isinstance(obj, torch.Tensor):
                return obj.to(self.device)
        except ImportError:
            pass
        if isinstance(obj, (tuple, list)):
            return type(obj)(self._to_device(x) for x in obj)
        if isinstance(obj, dict):
            return {k: self._to_device(v) for k, v in obj.items()}
        return obj

    def _safe_op(self, tensor: Any, op: Callable, op_name: str) -> Any:
        try:
            return op(tensor)
        except RuntimeError as e:
            err_str = str(e)
            if "MPS" in err_str or "not implemented" in err_str or "not supported" in err_str:
                logger.warning(f"{op_name}: MPS failed, falling back to CPU")
                cpu_input = self._to_cpu(tensor)
                cpu_result = op(cpu_input)
                return self._to_device(cpu_result)
            raise

    def inpaint_with_downscale_fallback(
        self,
        frames: list[np.ndarray],
        masks: list[np.ndarray],
        scale: float = 0.5,
    ) -> list[np.ndarray]:
        try:
            return self.inpaint(frames, masks)
        except (RuntimeError, InpaintTimeoutError) as e:
            logger.warning(f"Inpainting failed ({e}), retrying at {scale}x scale")

        h, w = frames[0].shape[:2]
        new_h, new_w = int(h * scale), int(w * scale)

        scaled_frames = [cv2.resize(f, (new_w, new_h), interpolation=cv2.INTER_AREA) for f in frames]
        scaled_masks = [
            cv2.resize(m.astype(np.uint8), (new_w, new_h), interpolation=cv2.INTER_NEAREST).astype(bool)
            for m in masks
        ]

        try:
            import torch
            torch.mps.empty_cache()
        except Exception:
            pass

        result_scaled = self.inpaint(scaled_frames, scaled_masks)
        return [cv2.resize(f, (w, h), interpolation=cv2.INTER_CUBIC) for f in result_scaled]

    def release(self) -> None:
        self._raft_model = None
        self._flow_comp_model = None
        self._inpaint_model = None
        self._loaded = False

        try:
            import torch
            torch.mps.synchronize()
            torch.mps.empty_cache()
        except Exception:
            pass

        logger.info("ProPainter engine released")
