from pathlib import Path

import cv2
import numpy as np
import pytest

from orisys.frame_processors import FrameUnwarper
from orisys.unroll import UVDomain


REPO_ROOT = Path(__file__).resolve().parents[1]
UV_MAP_PATH = REPO_ROOT / "unroll_map_uv.npz"


@pytest.fixture
def uv_domain():
    if not UV_MAP_PATH.is_file():
        pytest.skip(f"UV remap cache not found: {UV_MAP_PATH}")
    return UVDomain.load(str(UV_MAP_PATH))


def test_uv_unroll_processor_output_shape(uv_domain):
    uv_remap = uv_domain.remap
    src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
    src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))

    processor = FrameUnwarper(
        uv_remap,
        source_width=src_w,
        source_height=src_h,
        flip_source_vertical=True,
        border_value=(0, 0, 0),
    )

    frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
    unrolled = processor.process(frame)

    assert unrolled.shape == (uv_remap.out_height, uv_remap.out_width, 3)
    assert processor.valid_mask.shape == (uv_remap.out_height, uv_remap.out_width)
    assert unrolled.dtype == np.uint8


def test_uv_unroll_processor_prepare_source(uv_domain):
    uv_remap = uv_domain.remap
    src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
    src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))

    processor = FrameUnwarper(
        uv_remap,
        source_width=src_w,
        source_height=src_h,
        flip_source_vertical=False,
    )
    frame = np.ones((src_h, src_w, 3), dtype=np.uint8) * 128
    prepared = processor.prepare_source(frame)
    assert prepared.shape == (src_h, src_w, 3)


def test_uv_unroll_processor_flatten_prepared_matches_process(uv_domain):
    uv_remap = uv_domain.remap
    src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
    src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))

    processor = FrameUnwarper(
        uv_remap,
        source_width=src_w,
        source_height=src_h,
        flip_source_vertical=True,
        border_value=(0, 0, 0),
    )
    frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
    prepared = processor.prepare_source(frame)
    assert np.array_equal(processor.flatten_prepared(prepared), processor.process(frame))


def test_uv_unroll_processor_gray_fast_path_matches_bgr_result(uv_domain):
    uv_remap = uv_domain.remap
    src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
    src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))

    processor = FrameUnwarper(
        uv_remap,
        source_width=src_w,
        source_height=src_h,
        flip_source_vertical=True,
        border_value=(0, 0, 0),
    )
    frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
    prepared_gray = processor.prepare_source_gray(frame)
    gray_unrolled = processor.flatten_prepared_gray(prepared_gray)
    bgr_unrolled = processor.process(frame)
    assert gray_unrolled.shape == (uv_remap.out_height, uv_remap.out_width)
    assert np.allclose(gray_unrolled, cv2.cvtColor(bgr_unrolled, cv2.COLOR_BGR2GRAY), atol=1)