from __future__ import annotations

import json
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Tuple
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

import torch
from PIL import Image, UnidentifiedImageError
from torch.utils.data import DataLoader, Dataset


def load_json(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


class ChairAttributeDataset(Dataset):
    def __init__(
        self,
        split_path: Path,
        label_vocab_path: Path,
        repo_root: Path,
        image_cache_dir: Path,
    ) -> None:
        payload = load_json(split_path)
        if "samples" not in payload:
            raise RuntimeError(f"Invalid split file format: {split_path}")
        self.samples: List[Dict[str, Any]] = payload["samples"]
        self.label_vocab: Dict[str, int] = load_json(label_vocab_path)
        self.num_labels = len(self.label_vocab)
        self.repo_root = repo_root
        self.image_cache_dir = image_cache_dir
        self.image_cache_dir.mkdir(parents=True, exist_ok=True)
        self.corrupt_image_count = 0

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

    def _resolve_image_path(self, row: Dict[str, Any]) -> Path:
        local_rel = row.get("image_path", "")
        if local_rel:
            local = self.repo_root / local_rel
            if local.exists():
                return local
        image_id = str(row.get("image_id"))
        return self.image_cache_dir / f"{image_id}.jpg"

    def _candidate_urls(self, row: Dict[str, Any]) -> List[str]:
        primary = str(row.get("image_url") or "").strip()
        if not primary:
            return []
        candidates = [primary]
        if "/VG_100K/" in primary:
            candidates.append(primary.replace("/VG_100K/", "/VG_100K_2/"))
        elif "/VG_100K_2/" in primary:
            candidates.append(primary.replace("/VG_100K_2/", "/VG_100K/"))
        deduped: List[str] = []
        seen = set()
        for url in candidates:
            if url in seen:
                continue
            seen.add(url)
            deduped.append(url)
        return deduped

    def _download_if_missing(self, row: Dict[str, Any], path: Path, force: bool = False) -> None:
        if path.exists() and not force:
            return
        urls = self._candidate_urls(row)
        if not urls:
            raise RuntimeError(f"Missing image_url for image_id={row.get('image_id')}")
        last_error: Exception | None = None
        for url in urls:
            try:
                request = Request(url, headers={"User-Agent": "Mozilla/5.0"})
                with urlopen(request, timeout=30) as response:
                    image_bytes = response.read()
                # Validate bytes are actually decodable image data before caching.
                with Image.open(BytesIO(image_bytes)) as probe:
                    probe.verify()
                path.parent.mkdir(parents=True, exist_ok=True)
                with path.open("wb") as handle:
                    handle.write(image_bytes)
                return
            except (HTTPError, URLError, TimeoutError, OSError, UnidentifiedImageError) as exc:
                last_error = exc
                continue
        raise RuntimeError(
            f"Failed to download valid image for image_id={row.get('image_id')} from candidate URLs."
        ) from last_error

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

    def _load_image(self, row: Dict[str, Any]) -> Image.Image:
        image_path = self._resolve_image_path(row)
        self._download_if_missing(row, image_path)
        try:
            return self._open_local_image(image_path)
        except (UnidentifiedImageError, OSError):
            self.corrupt_image_count += 1
            if image_path.exists():
                image_path.unlink(missing_ok=True)
            # Retry once by re-downloading. Some VG image URLs occasionally produce
            # partial/bad files; this makes loading resilient instead of crashing workers.
            try:
                self._download_if_missing(row, image_path, force=True)
                return self._open_local_image(image_path)
            except Exception:
                return Image.new("RGB", (224, 224), color=(0, 0, 0))

    def _labels_to_multihot(self, labels: List[str]) -> torch.Tensor:
        target = torch.zeros(self.num_labels, dtype=torch.float32)
        for label in labels:
            if label in self.label_vocab:
                target[self.label_vocab[label]] = 1.0
        return target

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        row = self.samples[idx]
        image = self._load_image(row)
        labels = row.get("attributes_norm", [])
        target = self._labels_to_multihot(labels)
        return {
            "image": image,
            "target": target,
            "image_id": int(row["image_id"]),
            "labels": labels,
        }


def collate_samples(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    return {
        "images": [item["image"] for item in batch],
        "targets": torch.stack([item["target"] for item in batch], dim=0),
        "image_ids": [item["image_id"] for item in batch],
        "labels": [item["labels"] for item in batch],
    }


def build_dataloader(
    dataset: Dataset,
    batch_size: int,
    shuffle: bool,
    num_workers: int,
    pin_memory: bool,
) -> DataLoader:
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        collate_fn=collate_samples,
    )


def class_pos_weights(label_frequencies: Dict[str, int], label_vocab: Dict[str, int]) -> torch.Tensor:
    counts = torch.ones(len(label_vocab), dtype=torch.float32)
    for label, idx in label_vocab.items():
        counts[idx] = float(max(1, label_frequencies.get(label, 1)))
    weights = counts.max() / counts
    return weights


def split_paths(processed_root: Path, split_dir_name: str = "splits") -> Dict[str, Path]:
    split_dir = processed_root / split_dir_name
    return {
        "train": split_dir / "train.json",
        "val": split_dir / "val.json",
        "test": split_dir / "test.json",
    }