import importlib.util
from pathlib import Path

import cv2
import numpy as np
import torch


MODULE_PATH = Path(__file__).resolve().parents[1] / "frame_selector.py"
SPEC = importlib.util.spec_from_file_location("frame_selector", MODULE_PATH)
frame_selector = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(frame_selector)

AIA_SelectVideoFrame = frame_selector.AIA_SelectVideoFrame
_normalize_selector = frame_selector._normalize_selector
_selector_rank = frame_selector._selector_rank


def _checker(size=64, block=8):
    y, x = np.indices((size, size))
    pattern = ((x // block + y // block) % 2) * 255
    return np.stack([pattern, pattern, pattern], axis=-1).astype(np.uint8)


def _to_tensor(frames):
    return torch.from_numpy(np.stack(frames).astype(np.float32) / 255.0)


def _disable_face_detection(monkeypatch, count):
    monkeypatch.setattr(
        AIA_SelectVideoFrame,
        "_detect_face_scores",
        staticmethod(lambda _frames: np.zeros(count, dtype=np.float32)),
    )


def test_selector_aliases():
    assert _normalize_selector("") == "best"
    assert _normalize_selector("best still") == "best"
    assert _normalize_selector("no face please") == "no_face"
    assert _normalize_selector("very blurry") == "blurry"
    assert _normalize_selector("moving") == "moving"
    assert _normalize_selector("static pose") == "static"
    assert _normalize_selector("strong lighting flicker") == "lighting"
    assert _normalize_selector("whatever") == "best"
    assert _selector_rank("moving") == 1
    assert _selector_rank("moving rank 2") == 2
    assert _selector_rank("lighting #3") == 3


def test_ranked_indices_spread_candidates_apart():
    metrics = {
        "blur": np.asarray([0, 10, 100, 99, 20, 10, 0, 94, 0, 93], dtype=np.float32),
        "motion": np.zeros(10, dtype=np.float32),
        "lighting": np.zeros(10, dtype=np.float32),
        "face": np.ones(10, dtype=np.float32),
    }

    indices = AIA_SelectVideoFrame._ranked_indices(
        metrics,
        "best",
        count=3,
        min_separation=3,
    )

    assert indices == [2, 7, 9]


def test_blurry_mode_selects_soft_frame(monkeypatch):
    sharp = _checker()
    soft = cv2.GaussianBlur(sharp, (17, 17), 0)
    frames = [sharp, soft, sharp]
    _disable_face_detection(monkeypatch, len(frames))

    selected = AIA_SelectVideoFrame().select(_to_tensor(frames), "blurry")[0]

    assert torch.allclose(selected[0], _to_tensor([soft])[0])


def test_moving_and_static_modes_split_on_motion(monkeypatch):
    base = _checker()
    shifted = np.roll(base, 18, axis=1)
    frames = [base, base, base, shifted, base, base, base]
    _disable_face_detection(monkeypatch, len(frames))
    node = AIA_SelectVideoFrame()

    moving = node.select(_to_tensor(frames), "moving")[0]
    static = node.select(_to_tensor(frames), "static")[0]

    assert torch.allclose(moving[0], _to_tensor([shifted])[0])
    assert torch.allclose(static[0], _to_tensor([base])[0])


def test_lighting_mode_selects_brightness_jump(monkeypatch):
    dark = np.full((64, 64, 3), 40, dtype=np.uint8)
    pre_jump = np.full((64, 64, 3), 80, dtype=np.uint8)
    bright = np.full((64, 64, 3), 220, dtype=np.uint8)
    post_jump = np.full((64, 64, 3), 90, dtype=np.uint8)
    frames = [dark.copy() for _ in range(31)]
    frames[15] = pre_jump
    frames[16] = bright
    frames[17] = post_jump
    _disable_face_detection(monkeypatch, len(frames))

    selected = AIA_SelectVideoFrame().select(_to_tensor(frames), "flicker")[0]

    selected_mean = int(round(float(selected[0].mean() * 255)))
    assert selected_mean in {80, 90, 220}


def test_no_face_mode_uses_face_scores(monkeypatch):
    frames = [_checker() for _ in range(5)]
    face_scores = np.asarray([1.0, 0.8, 0.0, 0.7, 1.0], dtype=np.float32)
    monkeypatch.setattr(
        AIA_SelectVideoFrame,
        "_detect_face_scores",
        staticmethod(lambda _frames: face_scores),
    )

    selected = AIA_SelectVideoFrame().select(_to_tensor(frames), "no face")[0]

    assert torch.allclose(selected[0], _to_tensor([frames[2]])[0])


def test_ranked_selector_can_pick_second_distinct_frame(monkeypatch):
    frames = [_checker() for _ in range(40)]
    frames[10] = np.full((64, 64, 3), 210, dtype=np.uint8)
    frames[30] = np.full((64, 64, 3), 180, dtype=np.uint8)
    metrics = {
        "blur": np.asarray([1.0] * 40, dtype=np.float32),
        "motion": np.asarray([0.0] * 40, dtype=np.float32),
        "lighting": np.asarray([0.0] * 40, dtype=np.float32),
        "face": np.asarray([1.0] * 40, dtype=np.float32),
    }
    metrics["motion"][10] = 10.0
    metrics["motion"][30] = 9.0
    monkeypatch.setattr(
        AIA_SelectVideoFrame,
        "_metrics",
        classmethod(lambda _cls, _frames: metrics),
    )

    selected = AIA_SelectVideoFrame().select(_to_tensor(frames), "moving rank 2")[0]

    assert torch.allclose(selected[0], _to_tensor([frames[30]])[0])


def test_smoke_241_frame_batch(monkeypatch):
    frames = []
    base = _checker()
    for idx in range(241):
        frame = np.roll(base, idx % 32, axis=1)
        frames.append(frame)
    _disable_face_detection(monkeypatch, len(frames))

    selected = AIA_SelectVideoFrame().select(_to_tensor(frames), "")[0]

    assert selected.shape == (1, 64, 64, 3)