"""Shared fixtures for the augmentation pipeline tests."""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pytest
import torch
from PIL import Image

# Ensure src package is importable.
_SRC_ROOT = Path(__file__).resolve().parent.parent
if str(_SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(_SRC_ROOT))

from src.conf.hardware_conf import HardwareConfig
from src.conf.input_conf import InputConfig
from src.conf.models_conf import ModelsConfig
from src.conf.pipeline_conf import PipelineConfig
from src.conf.seg_conf import SegConfig
from src.augmentor.dataset import ImageRecord


# ---------------------------------------------------------------------------
# Config fixtures (no gin, plain defaults)
# ---------------------------------------------------------------------------

@pytest.fixture()
def hw_conf() -> HardwareConfig:
    import gin
    gin.clear_config()
    return HardwareConfig(
        profile_name="test", total_ram_gb=8.0, reserve_gb=1.0,
        use_fp16=False, batch_size=2, num_workers=0,
    )


@pytest.fixture()
def input_conf() -> InputConfig:
    import gin
    gin.clear_config()
    return InputConfig(image_size=64)


@pytest.fixture()
def seg_conf() -> SegConfig:
    import gin
    gin.clear_config()
    return SegConfig(prompts=["background", "building", "road"])


@pytest.fixture()
def models_conf() -> ModelsConfig:
    import gin
    gin.clear_config()
    return ModelsConfig()


@pytest.fixture()
def pipeline_conf(tmp_path: Path) -> PipelineConfig:
    import gin
    gin.clear_config()
    return PipelineConfig(
        input_root=str(tmp_path / "input"),
        output_root=str(tmp_path / "output"),
        stages=["depth", "edges", "segmentation", "chmv2"],
        save_npy=True,
        save_vis=True,
        save_concat=False,
        resume=False,
        log_level="WARNING",
    )


# ---------------------------------------------------------------------------
# Synthetic data fixtures
# ---------------------------------------------------------------------------

@pytest.fixture()
def fake_image_dir(tmp_path: Path) -> Path:
    """Create a minimal World-UAV-like directory with 4 dummy images."""
    img_dir = tmp_path / "input" / "scene01" / "query"
    img_dir.mkdir(parents=True)
    for i in range(4):
        img = Image.fromarray(
            np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8),
        )
        img.save(img_dir / f"img_{i:03d}.png")
    return tmp_path / "input"


@pytest.fixture()
def sample_batch() -> torch.Tensor:
    """Random [B, 3, H, W] float32 [0, 1] batch."""
    return torch.rand(2, 3, 64, 64)


@pytest.fixture()
def sample_depth() -> torch.Tensor:
    """Random [B, 1, H, W] float32 [0, 1] depth batch."""
    return torch.rand(2, 1, 64, 64)


@pytest.fixture()
def sample_records(fake_image_dir: Path, tmp_path: Path) -> list[ImageRecord]:
    """ImageRecord list matching fake_image_dir."""
    from src.augmentor.dataset import discover_images, attach_output_dirs
    records = discover_images(fake_image_dir)
    return attach_output_dirs(records, tmp_path / "output")
