from __future__ import annotations

import json
import random
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable, List, Optional

import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision.transforms import functional as TF


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}


def _iter_images(root: Path) -> Iterable[Path]:
    if not root.exists():
        return
    for path in root.rglob("*"):
        if path.suffix.lower() in IMAGE_EXTENSIONS and path.is_file():
            yield path


def _resolve_full_stem(candidate_stem: str, full_stems: set[str]) -> Optional[str]:
    if candidate_stem in full_stems:
        return candidate_stem

    base, sep, suffix = candidate_stem.rpartition("_")
    if sep and suffix.isdigit() and base in full_stems:
        return base
    return None


def _assign_splits(
    records: List[dict],
    train_ratio: float,
    val_ratio: float,
    seed: int,
) -> None:
    grouped: Dict[int, List[dict]] = defaultdict(list)
    for record in records:
        grouped[record["artist_id"]].append(record)

    rng = random.Random(seed)
    for artist_records in grouped.values():
        rng.shuffle(artist_records)
        n = len(artist_records)
        if n <= 2:
            splits = ["train"] * n
        else:
            n_train = max(1, int(n * train_ratio))
            n_val = max(1, int(n * val_ratio))
            if n_train + n_val >= n:
                n_val = max(1, n - n_train - 1)
            if n_train + n_val >= n:
                n_train = max(1, n - n_val - 1)
            if n_train + n_val >= n:
                splits = ["train"] * n
            else:
                splits = ["train"] * n_train + ["val"] * n_val + ["test"] * (n - n_train - n_val)

        for record, split in zip(artist_records, splits):
            record["split"] = split


def build_manifest(
    full_root: str | Path,
    face_root: str | Path,
    eye_root: str | Path,
    output_path: Optional[str | Path] = None,
    train_ratio: float = 0.7,
    val_ratio: float = 0.15,
    seed: int = 42,
) -> List[dict]:
    full_root = Path(full_root).resolve()
    face_root = Path(face_root).resolve()
    eye_root = Path(eye_root).resolve()

    if not full_root.exists():
        raise FileNotFoundError(f"full_root not found: {full_root}")

    full_by_artist: Dict[str, List[dict]] = defaultdict(list)
    for path in _iter_images(full_root):
        artist = path.parent.name
        full_by_artist[artist].append(
            {
                "artist": artist,
                "full_stem": path.stem,
                "full_path": str(path),
            }
        )

    face_candidates: Dict[tuple[str, str], List[dict]] = defaultdict(list)
    eye_candidates: Dict[tuple[str, str], dict] = defaultdict(dict)

    full_stems_by_artist = {
        artist: {entry["full_stem"] for entry in entries}
        for artist, entries in full_by_artist.items()
    }

    for path in _iter_images(face_root):
        artist = path.parent.name
        full_stem = _resolve_full_stem(path.stem, full_stems_by_artist.get(artist, set()))
        if full_stem is None:
            continue
        face_candidates[(artist, full_stem)].append(
            {
                "face_key": path.stem,
                "face_path": str(path),
            }
        )

    for path in _iter_images(eye_root):
        artist = path.parent.name
        eye_stem = path.stem
        side = None
        if eye_stem.endswith("_left"):
            face_key = eye_stem[:-5]
            side = "left"
        elif eye_stem.endswith("_right"):
            face_key = eye_stem[:-6]
            side = "right"
        if side is None:
            continue

        full_stem = _resolve_full_stem(face_key, full_stems_by_artist.get(artist, set()))
        if full_stem is None:
            continue
        eye_candidates[(artist, face_key)][side] = str(path)

    artists = sorted(full_by_artist)
    artist_to_id = {artist: idx for idx, artist in enumerate(artists)}
    records: List[dict] = []

    for artist in artists:
        artist_entries = sorted(full_by_artist[artist], key=lambda item: item["full_path"])
        for entry in artist_entries:
            full_stem = entry["full_stem"]
            candidates = sorted(
                face_candidates.get((artist, full_stem), []),
                key=lambda item: (item["face_key"] != full_stem, item["face_key"]),
            )
            chosen_face = None
            chosen_eyes: Dict[str, Optional[str]] = {}
            if candidates:
                chosen_face = candidates[0]
                chosen_eyes = eye_candidates.get((artist, chosen_face["face_key"]), {})
                for candidate in candidates:
                    eye_pair = eye_candidates.get((artist, candidate["face_key"]), {})
                    if eye_pair.get("left") or eye_pair.get("right"):
                        chosen_face = candidate
                        chosen_eyes = eye_pair
                        break

            sample_id = f"{artist}/{full_stem}"
            records.append(
                {
                    "sample_id": sample_id,
                    "artist": artist,
                    "artist_id": artist_to_id[artist],
                    "full_path": entry["full_path"],
                    "face_path": chosen_face["face_path"] if chosen_face is not None else None,
                    "face_key": chosen_face["face_key"] if chosen_face is not None else None,
                    "eye_left_path": chosen_eyes.get("left"),
                    "eye_right_path": chosen_eyes.get("right"),
                    "has_face": chosen_face is not None,
                    "has_eye": bool(chosen_eyes.get("left") or chosen_eyes.get("right")),
                }
            )

    _assign_splits(records, train_ratio=train_ratio, val_ratio=val_ratio, seed=seed)

    if output_path is not None:
        output_path = Path(output_path).resolve()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with output_path.open("w", encoding="utf-8") as handle:
            for record in records:
                handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    return records


def load_manifest(path: str | Path) -> List[dict]:
    path = Path(path).resolve()
    records: List[dict] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            records.append(json.loads(line))
    return records


class ArtistStyleDataset(Dataset):
    def __init__(
        self,
        records: List[dict],
        split: Optional[str] = None,
        image_size: int = 256,
        full_image_size: Optional[int] = None,
        face_image_size: Optional[int] = None,
        eye_image_size: Optional[int] = None,
        training: bool = False,
        seed: int = 42,
    ) -> None:
        if split is None:
            self.records = list(records)
        else:
            self.records = [record for record in records if record.get("split") == split]
        self.split = split
        self.image_size = int(image_size)
        self.full_image_size = int(full_image_size or image_size)
        self.face_image_size = int(face_image_size or image_size)
        self.eye_image_size = int(eye_image_size or image_size)
        self.training = training
        self.rng = random.Random(seed)
        self.artist_to_id = {
            record["artist"]: int(record["artist_id"])
            for record in records
        }
        self.id_to_artist = {idx: artist for artist, idx in self.artist_to_id.items()}

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

    def _load_image(self, path: str) -> Image.Image:
        with Image.open(path) as image:
            return image.convert("RGB")

    def _blank_image(self, size: int) -> Image.Image:
        return Image.new("RGB", (size, size), (0, 0, 0))

    def _select_eye_path(self, record: dict) -> Optional[str]:
        direct_eye_path = record.get("eye_path")
        if direct_eye_path:
            return direct_eye_path
        left_path = record.get("eye_left_path")
        right_path = record.get("eye_right_path")
        if left_path and right_path:
            if self.training:
                return left_path if self.rng.random() < 0.5 else right_path
            return left_path
        if left_path:
            return left_path
        if right_path:
            return right_path
        return None

    def _load_optional_image(self, path: Optional[str], size: int) -> tuple[Image.Image, bool]:
        if path:
            return self._load_image(path), True
        return self._blank_image(size), False

    def _apply_transforms(self, full: Image.Image, face: Image.Image, eye: Image.Image) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if self.training and self.rng.random() < 0.5:
            full = TF.hflip(full)
            face = TF.hflip(face)
            eye = TF.hflip(eye)

        full = TF.resize(full, [self.full_image_size, self.full_image_size], interpolation=Image.Resampling.BICUBIC)
        face = TF.resize(face, [self.face_image_size, self.face_image_size], interpolation=Image.Resampling.BICUBIC)
        eye = TF.resize(eye, [self.eye_image_size, self.eye_image_size], interpolation=Image.Resampling.BICUBIC)

        return TF.to_tensor(full), TF.to_tensor(face), TF.to_tensor(eye)

    def __getitem__(self, index: int) -> dict:
        record = self.records[index]
        full = self._load_image(record["full_path"])
        face, has_face = self._load_optional_image(record.get("face_path"), self.face_image_size)
        eye, has_eye = self._load_optional_image(self._select_eye_path(record), self.eye_image_size)
        full_tensor, face_tensor, eye_tensor = self._apply_transforms(full, face, eye)
        return {
            "full": full_tensor,
            "face": face_tensor,
            "eye": eye_tensor,
            "view_mask": torch.tensor([1.0, float(has_face), float(has_eye)], dtype=torch.float32),
            "label": torch.tensor(int(record.get("train_label", record["artist_id"])), dtype=torch.long),
            "artist": record["artist"],
            "sample_id": record["sample_id"],
        }