"""Integration tests: run stages end-to-end with mock models on synthetic data."""

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import numpy as np
import torch

from src.augmentor.dataset import (
    AugmentDataset,
    ImageRecord,
    attach_output_dirs,
    discover_images,
)
from src.augmentor.inference import compute_edges_from_depth
from src.augmentor.io_utils import npy_path, vis_path, shutdown_io_pool
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


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_mock_depth_model_da3(H: int, W: int) -> MagicMock:
    model = MagicMock()

    def inference(image_list, process_res=None):
        depth = np.random.rand(len(image_list), H, W).astype(np.float32)
        return SimpleNamespace(depth=depth)

    model.inference = inference
    return model


def _make_mock_chmv2(H: int, W: int) -> tuple[MagicMock, MagicMock]:
    mock_param = torch.nn.Parameter(torch.zeros(1))
    model = MagicMock()
    model.parameters = lambda: iter([mock_param])

    def forward(pixel_values=None):
        B = pixel_values.shape[0]
        return SimpleNamespace(predicted_depth=torch.rand(B, H, W))

    model.side_effect = forward

    processor = MagicMock()

    def preprocess(images=None, return_tensors=None):
        B = len(images)
        return {"pixel_values": torch.rand(B, 3, H, W)}

    processor.side_effect = preprocess

    def post_process(outputs, target_sizes=None):
        B = outputs.predicted_depth.shape[0]
        return [
            {"predicted_depth": torch.rand(target_sizes[i][0], target_sizes[i][1])}
            for i in range(B)
        ]

    processor.post_process_depth_estimation = post_process
    return model, processor


def _make_mock_segformer(num_classes: int, H: int, W: int) -> tuple[MagicMock, dict]:
    mock_param = torch.nn.Parameter(torch.zeros(1))
    model = MagicMock()
    model.parameters = lambda: iter([mock_param])

    def forward(pixel_values=None):
        B = pixel_values.shape[0]
        logits = torch.randn(B, num_classes, H // 4, W // 4)
        return SimpleNamespace(logits=logits)

    model.side_effect = forward

    processor = SimpleNamespace(
        image_mean=[0.485, 0.456, 0.406],
        image_std=[0.229, 0.224, 0.225],
    )
    seg_config = {"type": "segformer", "processor": processor, "num_classes": num_classes}
    return model, seg_config


# ---------------------------------------------------------------------------
# Per-stage tests
# ---------------------------------------------------------------------------

class TestDepthStageIntegration:
    def test_produces_depth_files(
        self,
        sample_records: list[ImageRecord],
        pipeline_conf: PipelineConfig,
        hw_conf: HardwareConfig,
        input_conf: InputConfig,
    ) -> None:
        from src.main import run_depth_stage

        mock_model = _make_mock_depth_model_da3(input_conf.image_size, input_conf.image_size)

        with patch("src.main.load_depth_model", return_value=mock_model), \
             patch("src.main.unload_model"):
            run_depth_stage(
                sample_records, pipeline_conf, hw_conf,
                ModelsConfig(), input_conf, torch.device("cpu"),
            )

        for r in sample_records:
            p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
            assert p.exists(), f"Missing depth npy for {r.stem}"
            arr = np.load(p)
            assert arr.dtype == np.float16
            assert arr.shape[0] == 1


class TestChmv2StageIntegration:
    def test_produces_chm_files(
        self,
        sample_records: list[ImageRecord],
        pipeline_conf: PipelineConfig,
        hw_conf: HardwareConfig,
        input_conf: InputConfig,
    ) -> None:
        from src.main import run_chmv2_stage

        mock_model, mock_processor = _make_mock_chmv2(
            input_conf.image_size, input_conf.image_size,
        )

        with patch("src.main.load_chmv2_model", return_value=(mock_model, mock_processor)), \
             patch("src.main.unload_model"):
            run_chmv2_stage(
                sample_records, pipeline_conf, hw_conf,
                ModelsConfig(), input_conf, torch.device("cpu"),
            )

        for r in sample_records:
            p = npy_path(r.output_root, "chm", r.rel_parent, r.stem)
            assert p.exists(), f"Missing chm npy for {r.stem}"
            arr = np.load(p)
            assert arr.dtype == np.float16
            assert arr.shape[0] == 1


class TestEdgesStageIntegration:
    def test_produces_edges_files(
        self,
        sample_records: list[ImageRecord],
        pipeline_conf: PipelineConfig,
    ) -> None:
        from src.main import run_edges_stage

        for r in sample_records:
            p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
            p.parent.mkdir(parents=True, exist_ok=True)
            np.save(p, np.random.rand(1, 64, 64).astype(np.float16))

        run_edges_stage(sample_records, pipeline_conf)

        for r in sample_records:
            p = npy_path(r.output_root, "edge", r.rel_parent, r.stem)
            assert p.exists(), f"Missing edge npy for {r.stem}"
            arr = np.load(p)
            assert arr.dtype == np.float16

    def test_skips_missing_depth(
        self,
        sample_records: list[ImageRecord],
        pipeline_conf: PipelineConfig,
    ) -> None:
        from src.main import run_edges_stage

        run_edges_stage(sample_records, pipeline_conf)
        for r in sample_records:
            assert not npy_path(r.output_root, "edge", r.rel_parent, r.stem).exists()


class TestSegmentationStageIntegration:
    def test_segformer_produces_seg_files(
        self,
        sample_records: list[ImageRecord],
        pipeline_conf: PipelineConfig,
        hw_conf: HardwareConfig,
        input_conf: InputConfig,
        seg_conf: SegConfig,
    ) -> None:
        from src.main import run_segmentation_stage

        num_classes = seg_conf.num_classes
        mock_model, seg_config_dict = _make_mock_segformer(
            num_classes, input_conf.image_size, input_conf.image_size,
        )

        with patch(
            "src.main.load_segmentation_model",
            return_value=(mock_model, seg_config_dict),
        ), patch("src.main.unload_model"):
            run_segmentation_stage(
                sample_records, pipeline_conf, hw_conf,
                ModelsConfig(), input_conf, seg_conf, torch.device("cpu"),
            )

        for r in sample_records:
            p = npy_path(r.output_root, "segm", r.rel_parent, r.stem)
            assert p.exists()
            arr = np.load(p)
            assert arr.dtype == np.uint8


# ---------------------------------------------------------------------------
# Full pipeline smoke test
# ---------------------------------------------------------------------------

class TestFullPipelineSmoke:
    def test_pipeline_completes(
        self,
        fake_image_dir: Path,
        tmp_path: Path,
    ) -> None:
        import gin
        gin.clear_config()

        pipeline_conf = PipelineConfig(
            input_root=str(fake_image_dir),
            output_root=str(tmp_path / "output"),
            stages=["depth", "edges", "segmentation", "chmv2"],
            save_npy=True,
            save_vis=False,
            save_concat=False,
            resume=False,
            log_level="WARNING",
        )
        hw_conf = HardwareConfig(
            profile_name="test", total_ram_gb=8.0, reserve_gb=1.0,
            batch_size=2, num_workers=0,
        )
        input_conf = InputConfig(image_size=32)
        seg_conf = SegConfig(prompts=["bg", "building", "road"])
        models_conf = ModelsConfig()

        H = W = input_conf.image_size
        num_classes = seg_conf.num_classes

        mock_depth = _make_mock_depth_model_da3(H, W)
        mock_chmv2_model, mock_chmv2_proc = _make_mock_chmv2(H, W)
        mock_seg, seg_config_dict = _make_mock_segformer(num_classes, H, W)

        with patch("src.main.load_depth_model", return_value=mock_depth), \
             patch("src.main.load_chmv2_model", return_value=(mock_chmv2_model, mock_chmv2_proc)), \
             patch("src.main.load_segmentation_model", return_value=(mock_seg, seg_config_dict)), \
             patch("src.main.unload_model"):
            from src.main import run_pipeline
            run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf)

        output_root = Path(pipeline_conf.output_root)
        assert (output_root / "manifest.json").exists()

        # New dir layout: npy/depth/..., npy/edge/..., etc.
        found_depth = list((output_root / "npy" / "depth").rglob("*.npy"))
        found_edges = list((output_root / "npy" / "edge").rglob("*.npy"))
        found_seg = list((output_root / "npy" / "segm").rglob("*.npy"))
        found_chmv2 = list((output_root / "npy" / "chm").rglob("*.npy"))
        assert len(found_depth) > 0
        assert len(found_edges) > 0
        assert len(found_seg) > 0
        assert len(found_chmv2) > 0


class TestFullPipelineVisOnly:
    def test_produces_png_not_npy(
        self, fake_image_dir: Path, tmp_path: Path,
    ) -> None:
        import gin
        gin.clear_config()

        pipeline_conf = PipelineConfig(
            input_root=str(fake_image_dir),
            output_root=str(tmp_path / "output"),
            stages=["depth", "edges", "segmentation", "chmv2"],
            save_npy=False, save_vis=True,
            save_safetensors=False,
            save_concat=False, resume=False, log_level="WARNING",
        )
        hw_conf = HardwareConfig(
            profile_name="test", total_ram_gb=8.0, reserve_gb=1.0,
            batch_size=2, num_workers=0,
        )
        input_conf = InputConfig(image_size=32)
        seg_conf = SegConfig(prompts=["bg", "building", "road"])
        models_conf = ModelsConfig()
        H = W = 32
        num_classes = seg_conf.num_classes

        mock_depth = _make_mock_depth_model_da3(H, W)
        mock_chmv2_model, mock_chmv2_proc = _make_mock_chmv2(H, W)
        mock_seg, seg_config_dict = _make_mock_segformer(num_classes, H, W)

        with patch("src.main.load_depth_model", return_value=mock_depth), \
             patch("src.main.load_chmv2_model", return_value=(mock_chmv2_model, mock_chmv2_proc)), \
             patch("src.main.load_segmentation_model", return_value=(mock_seg, seg_config_dict)), \
             patch("src.main.unload_model"):
            from src.main import run_pipeline
            run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf)

        output_root = Path(pipeline_conf.output_root)
        assert len(list((output_root / "depth").rglob("*.png"))) > 0
        assert len(list((output_root / "edge").rglob("*.png"))) > 0
        assert len(list((output_root / "segm").rglob("*.png"))) > 0
        assert len(list((output_root / "chm").rglob("*.png"))) > 0
        assert len(list(output_root.rglob("*.npy"))) == 0
