"""TuSimple (mask-based) dataset loader.

The data on disk (`<root>/training/{frames,lane-masks}/`) is a segmentation-style
export rather than the official `label_data_*.json` point annotations, and has no
official val/test split (see project plan §3.2). This loader:
  1. Pairs each frame with its mask by filename.
  2. Extracts per-lane polylines from the mask via `utils.mask_to_lanes`.
  3. Applies the letterbox/normalize/(flip) transform pipeline to produce
     fixed-size training targets.
  4. Carves a seeded train/val split from the single available folder.
"""
from __future__ import annotations

import os
import random
from pathlib import Path

import cv2
import torch
from torch.utils.data import Dataset

from utils.mask_to_lanes import mask_to_lanes
from data.transforms import build_transforms


class TuSimpleMaskDataset(Dataset):
    def __init__(
        self,
        root: str,
        split: str,
        out_w: int,
        out_h: int,
        max_lanes: int,
        num_sample_ys: int,
        val_fraction: float = 0.1,
        seed: int = 42,
    ):
        assert split in ("train", "val")
        self.root = Path(root)
        frames_dir = self.root / "training" / "frames"
        masks_dir = self.root / "training" / "lane-masks"

        frame_files = sorted(f for f in os.listdir(frames_dir) if f.lower().endswith((".jpg", ".jpeg", ".png")))
        pairs = [f for f in frame_files if (masks_dir / f).exists()]
        if not pairs:
            raise RuntimeError(f"No matching frame/mask pairs found under {self.root}")

        rng = random.Random(seed)
        shuffled = pairs[:]
        rng.shuffle(shuffled)
        n_val = max(1, int(len(shuffled) * val_fraction))
        val_set = set(shuffled[:n_val])

        if split == "val":
            self.files = [f for f in pairs if f in val_set]
        else:
            self.files = [f for f in pairs if f not in val_set]

        self.frames_dir = frames_dir
        self.masks_dir = masks_dir
        self.transforms = build_transforms(out_w, out_h, max_lanes, num_sample_ys, train=(split == "train"))

    def __len__(self) -> int:
        return len(self.files)

    def __getitem__(self, idx: int) -> dict:
        fname = self.files[idx]
        image_bgr = cv2.imread(str(self.frames_dir / fname))
        mask = cv2.imread(str(self.masks_dir / fname))
        if image_bgr is None or mask is None:
            raise RuntimeError(f"Failed to read {fname} from {self.frames_dir} / {self.masks_dir}")

        image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        lanes = mask_to_lanes(mask)

        sample = {"image": image, "lanes": lanes, "name": fname}
        sample = self.transforms(sample)
        return sample


def collate_fn(batch: list[dict]) -> dict:
    images = torch.stack([b["image"] for b in batch], dim=0)
    target_xs = torch.stack([b["target_xs"] for b in batch], dim=0)
    target_valid_mask = torch.stack([b["target_valid_mask"] for b in batch], dim=0)
    target_lane_valid = torch.stack([b["target_lane_valid"] for b in batch], dim=0)
    target_endpoints = torch.stack([b["target_endpoints"] for b in batch], dim=0)
    sample_ys = batch[0]["sample_ys"]
    names = [b["name"] for b in batch]
    return {
        "images": images,
        "target_xs": target_xs,
        "target_valid_mask": target_valid_mask,
        "target_lane_valid": target_lane_valid,
        "target_endpoints": target_endpoints,
        "sample_ys": sample_ys,
        "names": names,
    }