"""Quality metrics for inpainting verification (S2-R1-016).

Replaces PSNR/SSIM as repair quality metrics. Uses:
- Temporal consistency: adjacent frame L2 distance stddev in mask region
- Style consistency: LPIPS between repaired region and surrounding area
- Pixel integrity: PSNR of non-masked region (should be >50dB = untouched)
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from loguru import logger


@dataclass
class QualityReport:
    temporal_consistency: float   # lower is better (stddev of inter-frame diff)
    pixel_integrity_psnr: float   # higher is better (non-mask region, should be >50)
    style_lpips: float | None     # lower is better (None if lpips unavailable)
    passed: bool
    details: dict


def compute_temporal_consistency(
    frames: list[np.ndarray],
    masks: list[np.ndarray],
) -> float:
    """Compute temporal flicker score in masked region.

    Returns mean of per-frame L2 differences in the inpainted area.
    Lower = more temporally consistent (less flicker).
    Constant content → 0, alternating black/white → ~255.
    """
    if len(frames) < 2:
        return 0.0

    diffs: list[float] = []
    for i in range(1, len(frames)):
        mask = masks[i] if i < len(masks) else masks[-1]
        if mask.sum() == 0:
            continue

        diff = np.abs(frames[i].astype(np.float32) - frames[i - 1].astype(np.float32))
        masked_diff = diff[mask]
        if masked_diff.size > 0:
            diffs.append(float(masked_diff.mean()))

    if not diffs:
        return 0.0

    return float(np.mean(diffs))


def compute_pixel_integrity(
    original_frames: list[np.ndarray],
    output_frames: list[np.ndarray],
    masks: list[np.ndarray],
) -> float:
    """Compute PSNR of non-masked region between original and output.

    Should be very high (>50dB) if non-mask pixels are untouched.
    """
    if not original_frames or not output_frames:
        return 0.0

    n = min(len(original_frames), len(output_frames))
    mse_sum = 0.0
    pixel_count = 0

    for i in range(n):
        mask = masks[i] if i < len(masks) else masks[-1]
        non_mask = ~mask

        if non_mask.sum() == 0:
            continue

        orig = original_frames[i].astype(np.float64)
        out = output_frames[i].astype(np.float64)

        diff_sq = (orig[non_mask] - out[non_mask]) ** 2
        mse_sum += diff_sq.sum()
        pixel_count += diff_sq.size

    if pixel_count == 0:
        return float("inf")

    mse = mse_sum / pixel_count
    if mse < 1e-10:
        return float("inf")

    psnr = 10 * np.log10(255.0 ** 2 / mse)
    return float(psnr)


def compute_style_lpips(
    frames: list[np.ndarray],
    masks: list[np.ndarray],
    sample_count: int = 10,
) -> float | None:
    """Compute LPIPS between repaired region and surrounding context.

    Returns None if lpips package is not available.
    """
    try:
        import torch
        import lpips
    except ImportError:
        logger.warning("lpips not installed, skipping style consistency check")
        return None

    loss_fn = lpips.LPIPS(net="alex", verbose=False)

    scores: list[float] = []
    step = max(1, len(frames) // sample_count)

    for i in range(0, len(frames), step):
        mask = masks[i] if i < len(masks) else masks[-1]
        if mask.sum() == 0:
            continue

        frame = frames[i]
        h, w = frame.shape[:2]

        # Get bounding box of mask
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)
        if not rows.any() or not cols.any():
            continue

        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]

        # Expand region for context (2x the mask bbox)
        pad_h = (rmax - rmin) // 2
        pad_w = (cmax - cmin) // 2
        ctx_rmin = max(0, rmin - pad_h)
        ctx_rmax = min(h, rmax + pad_h)
        ctx_cmin = max(0, cmin - pad_w)
        ctx_cmax = min(w, cmax + pad_w)

        patch = frame[ctx_rmin:ctx_rmax, ctx_cmin:ctx_cmax]
        if patch.shape[0] < 8 or patch.shape[1] < 8:
            continue

        # Compare mask region vs context region using LPIPS
        # We create two patches: one from mask area, one from surrounding
        mask_crop = mask[ctx_rmin:ctx_rmax, ctx_cmin:ctx_cmax]

        # Masked pixels patch
        masked_patch = patch.copy()
        masked_patch[~mask_crop] = 0

        # Context pixels patch
        ctx_patch = patch.copy()
        ctx_patch[mask_crop] = 0

        # Convert to torch
        t1 = torch.from_numpy(masked_patch).permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1
        t2 = torch.from_numpy(ctx_patch).permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1

        with torch.no_grad():
            score = loss_fn(t1, t2).item()
        scores.append(score)

    if not scores:
        return None

    return float(np.mean(scores))


def run_quality_check(
    original_frames: list[np.ndarray],
    output_frames: list[np.ndarray],
    masks: list[np.ndarray],
    temporal_threshold: float = 5.0,
    integrity_threshold: float = 50.0,
    lpips_threshold: float = 0.1,
) -> QualityReport:
    """Run all quality checks and return report."""
    temporal = compute_temporal_consistency(output_frames, masks)
    integrity = compute_pixel_integrity(original_frames, output_frames, masks)
    style = compute_style_lpips(output_frames, masks)

    passed = True
    details: dict = {}

    if temporal > temporal_threshold:
        passed = False
        details["temporal"] = f"FAIL: {temporal:.2f} > {temporal_threshold}"
    else:
        details["temporal"] = f"PASS: {temporal:.2f} <= {temporal_threshold}"

    if integrity < integrity_threshold:
        passed = False
        details["integrity"] = f"FAIL: {integrity:.1f}dB < {integrity_threshold}dB"
    else:
        details["integrity"] = f"PASS: {integrity:.1f}dB >= {integrity_threshold}dB"

    if style is not None:
        if style > lpips_threshold:
            details["style"] = f"WARN: LPIPS {style:.3f} > {lpips_threshold}"
        else:
            details["style"] = f"PASS: LPIPS {style:.3f} <= {lpips_threshold}"
    else:
        details["style"] = "SKIP: lpips not available"

    return QualityReport(
        temporal_consistency=temporal,
        pixel_integrity_psnr=integrity,
        style_lpips=style,
        passed=passed,
        details=details,
    )
