from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING, Any

from .keypoints import normalize_keypoints

if TYPE_CHECKING:  # pragma: no cover
    import numpy as np


def _first_key(data: dict[str, Any], keys: tuple[str, ...]) -> Any:
    for key in keys:
        if key in data:
            return data[key]
    return None


def _resolve_optional_path(value: Any, root: Path) -> Any:
    if value is None:
        return None
    if isinstance(value, (list, tuple)):
        if not value:
            return list(value)
        if isinstance(value[0], (str, Path)):
            out = []
            for item in value:
                if item is None:
                    out.append(None)
                    continue
                path = Path(item)
                if not path.is_absolute():
                    path = root / path
                out.append(str(path))
            return out
        return value
    if isinstance(value, str) or isinstance(value, Path):
        path = Path(value)
        if not path.is_absolute():
            path = root / path
        return str(path)
    return value


def _load_sidecar_metadata(meta_path: Path, root: Path) -> dict[str, Any]:
    """Load optional per-image metadata from labels/<split>/<image>.json.

    This loader is intentionally lightweight: it resolves relative paths but does
    not decode masks/depth arrays. The resulting keys are merged into the record.
    """

    if not meta_path.exists():
        return {}
    try:
        data = json.loads(meta_path.read_text())
    except Exception:
        return {}

    out: dict[str, Any] = {}

    # Optional image size hints (used for pixel conversion).
    image_hw = data.get("image_hw") or data.get("hw")
    image_size = data.get("image_size")
    if image_hw is not None:
        out["image_hw"] = image_hw
    if image_size is not None:
        out["image_size"] = image_size

    # Pose / intrinsics.
    if "pose" in data:
        out["pose"] = data.get("pose")
    r_gt = _first_key(data, ("R_gt", "R"))
    t_gt = _first_key(data, ("t_gt", "t"))
    if r_gt is not None:
        out["R_gt"] = r_gt
    if t_gt is not None:
        out["t_gt"] = t_gt
    if "pose" not in out and (r_gt is not None or t_gt is not None):
        out["pose"] = {"R": r_gt, "t": t_gt}

    if "intrinsics" in data:
        out["intrinsics"] = data.get("intrinsics")
    k_gt = _first_key(data, ("K_gt", "K"))
    if k_gt is not None:
        out["K_gt"] = k_gt
        out["intrinsics"] = k_gt

    # Optional paths (kept as paths/inline values).
    mask_value = _first_key(data, ("mask_path", "M_path", "M", "mask"))
    depth_value = _first_key(data, ("depth_path", "D_obj_path", "D_obj", "depth"))
    # Optional mask metadata (instance/semantic segmentation helpers).
    if "mask_format" in data:
        out["mask_format"] = data.get("mask_format")
    if "mask_instances" in data:
        out["mask_instances"] = data.get("mask_instances")
    if "mask_classes" in data:
        out["mask_classes"] = data.get("mask_classes")
    if "mask_class_id" in data:
        out["mask_class_id"] = data.get("mask_class_id")
    if "mask_class_map" in data:
        out["mask_class_map"] = data.get("mask_class_map")
    if mask_value is not None:
        resolved = _resolve_optional_path(mask_value, root)
        out["mask_path"] = resolved
        out["mask"] = resolved
    if depth_value is not None:
        resolved = _resolve_optional_path(depth_value, root)
        out["depth_path"] = resolved
        out["depth"] = resolved
        out["D_obj"] = resolved

    return out


def _as_numpy_mask(mask: Any) -> np.ndarray:
    import numpy as np

    if isinstance(mask, np.ndarray):
        return mask
    return np.asarray(mask)


def _load_mask_value(value: Any) -> Any:
    import numpy as np
    from PIL import Image

    if value is None:
        return None
    if isinstance(value, (str, Path)):
        path = Path(value)
        if not path.exists():
            return None
        try:
            with Image.open(path) as img:
                return np.asarray(img)
        except Exception:
            return None
    return value


def _parse_color_key(key: Any) -> Any:
    if isinstance(key, (int, float)):
        return int(key)
    if isinstance(key, str):
        if "," in key:
            parts = [int(p.strip()) for p in key.split(",")]
            return tuple(parts)
        if key.startswith("#") and len(key) == 7:
            r = int(key[1:3], 16)
            g = int(key[3:5], 16)
            b = int(key[5:7], 16)
            return (r, g, b)
        if key.isdigit():
            return int(key)
    return key


def _bbox_from_mask(mask_bool: np.ndarray) -> tuple[int, int, int, int] | None:
    import numpy as np

    ys, xs = np.where(mask_bool)
    if ys.size == 0:
        return None
    y_min = int(ys.min())
    y_max = int(ys.max())
    x_min = int(xs.min())
    x_max = int(xs.max())
    return x_min, y_min, x_max, y_max


def _connected_components(mask_bool: np.ndarray) -> list[tuple[int, int, int, int]]:
    import numpy as np

    h, w = mask_bool.shape
    visited = np.zeros((h, w), dtype=bool)
    bboxes: list[tuple[int, int, int, int]] = []
    for y in range(h):
        for x in range(w):
            if not mask_bool[y, x] or visited[y, x]:
                continue
            stack = [(y, x)]
            visited[y, x] = True
            y_min = y_max = y
            x_min = x_max = x
            while stack:
                cy, cx = stack.pop()
                y_min = min(y_min, cy)
                y_max = max(y_max, cy)
                x_min = min(x_min, cx)
                x_max = max(x_max, cx)
                for ny, nx in ((cy - 1, cx), (cy + 1, cx), (cy, cx - 1), (cy, cx + 1)):
                    if 0 <= ny < h and 0 <= nx < w and mask_bool[ny, nx] and not visited[ny, nx]:
                        visited[ny, nx] = True
                        stack.append((ny, nx))
            bboxes.append((x_min, y_min, x_max, y_max))
    return bboxes


def _label_from_bbox(*, class_id: int, bbox: tuple[int, int, int, int], h: int, w: int) -> dict[str, Any]:
    x_min, y_min, x_max, y_max = bbox
    cx = (x_min + x_max + 1) / 2.0 / w
    cy = (y_min + y_max + 1) / 2.0 / h
    bw = (x_max - x_min + 1) / w
    bh = (y_max - y_min + 1) / h
    return {"class_id": int(class_id), "cx": float(cx), "cy": float(cy), "w": float(bw), "h": float(bh)}


def _mask_to_yolo_labels(
    mask: Any,
    *,
    mask_format: str | None = None,
    class_map: dict[str, Any] | None = None,
    mask_class_id: int | None = None,
    instances: bool = False,
) -> list[dict[str, Any]]:
    """Derive YOLO-style bbox labels from a mask image or array.

    Supported patterns:
    - color mask (RGB): each unique non-black color is treated as a class
      - optional mask_class_map maps colors (\"r,g,b\" or \"#RRGGBB\") -> class_id
    - instance mask (single-channel IDs): each non-zero ID is treated as an instance
      - class_id comes from mask_class_id (fallback) or mask_class_map (id -> class_id)
    - class-id mask (single-channel IDs): each non-zero value is treated as a class_id (or remapped via class_map)

    When instances=True, connected components are split into multiple boxes (useful for binary masks).
    """

    class_map = class_map or {}
    import numpy as np

    mask_np = _as_numpy_mask(mask)

    # RGB / RGBA color masks.
    if mask_format == "color" or (mask_np.ndim == 3 and mask_np.shape[-1] in (3, 4)):
        if mask_np.ndim != 3:
            return []
        rgb = mask_np[..., :3]
        flat = rgb.reshape(-1, 3)
        unique = np.unique(flat, axis=0)
        colors = [tuple(c.tolist()) for c in unique]
        colors = [c for c in colors if c != (0, 0, 0)]
        mapping: dict[tuple[int, int, int], int] = {}
        if class_map:
            for k, v in class_map.items():
                key = _parse_color_key(k)
                if isinstance(key, tuple) and len(key) == 3:
                    mapping[tuple(int(x) for x in key)] = int(v)
        else:
            for idx, color in enumerate(sorted(colors)):
                mapping[color] = int(idx)

        labels: list[dict[str, Any]] = []
        h, w = mask_np.shape[:2]
        for color, class_id in mapping.items():
            mask_bool = np.all(rgb == np.array(color, dtype=rgb.dtype), axis=-1)
            if instances:
                for bbox in _connected_components(mask_bool):
                    labels.append(_label_from_bbox(class_id=int(class_id), bbox=bbox, h=h, w=w))
            else:
                bbox = _bbox_from_mask(mask_bool)
                if bbox is None:
                    continue
                labels.append(_label_from_bbox(class_id=int(class_id), bbox=bbox, h=h, w=w))
        return labels

    # Numeric masks: instance IDs or class IDs.
    if mask_np.ndim == 3 and mask_np.shape[-1] == 1:
        mask_np = mask_np[..., 0]
    if mask_np.ndim != 2:
        return []
    mask_int = mask_np.astype(int)
    unique = np.unique(mask_int)
    ids = [int(v) for v in unique if int(v) != 0]
    h, w = mask_int.shape[:2]

    labels = []
    fmt = (mask_format or "").lower()
    is_instance_mask = fmt in ("instance", "instances")
    fallback_class = int(mask_class_id) if mask_class_id is not None else 0
    for mask_id in ids:
        class_id = fallback_class
        if class_map:
            class_id = int(class_map.get(str(mask_id), class_map.get(mask_id, class_id)))
        elif not is_instance_mask:
            class_id = int(mask_id)
        mask_bool = mask_int == int(mask_id)
        if instances and not is_instance_mask:
            for bbox in _connected_components(mask_bool):
                labels.append(_label_from_bbox(class_id=int(class_id), bbox=bbox, h=h, w=w))
            continue
        bbox = _bbox_from_mask(mask_bool)
        if bbox is None:
            continue
        labels.append(_label_from_bbox(class_id=int(class_id), bbox=bbox, h=h, w=w))
    return labels


def _derive_labels_from_mask(record: dict[str, Any]) -> list[dict[str, Any]]:
    if record.get("mask_path") is None and record.get("mask") is None:
        return []

    import numpy as np

    mask_value = record.get("mask") or record.get("mask_path")
    instances = bool(record.get("mask_instances", False))
    class_map = record.get("mask_class_map") or {}
    mask_classes = record.get("mask_classes")
    mask_format = record.get("mask_format")
    mask_class_id = record.get("mask_class_id")

    if isinstance(mask_value, (list, tuple)) and mask_value and isinstance(mask_value[0], (str, Path)):
        labels: list[dict[str, Any]] = []
        class_ids = list(mask_classes) if isinstance(mask_classes, (list, tuple)) else list(range(len(mask_value)))
        for idx, mask_path in enumerate(mask_value):
            mask_data = _load_mask_value(mask_path)
            if mask_data is None:
                continue
            class_id = class_ids[idx] if idx < len(class_ids) else idx
            mask_np = _as_numpy_mask(mask_data)
            if mask_np.ndim == 3 and mask_np.shape[-1] in (3, 4):
                mask_bool = np.any(mask_np[..., :3] != 0, axis=-1)
            elif mask_np.ndim == 3 and mask_np.shape[-1] == 1:
                mask_bool = mask_np[..., 0] != 0
            else:
                mask_bool = mask_np != 0
            if mask_bool.ndim != 2:
                continue
            h, w = mask_bool.shape
            if instances:
                for bbox in _connected_components(mask_bool):
                    labels.append(_label_from_bbox(class_id=int(class_id), bbox=bbox, h=h, w=w))
            else:
                bbox = _bbox_from_mask(mask_bool)
                if bbox is None:
                    continue
                labels.append(_label_from_bbox(class_id=int(class_id), bbox=bbox, h=h, w=w))
        return labels

    mask_data = _load_mask_value(mask_value)
    if mask_data is None:
        return []
    return _mask_to_yolo_labels(
        mask_data,
        mask_format=str(mask_format) if mask_format is not None else None,
        class_map=class_map if isinstance(class_map, dict) else {},
        mask_class_id=int(mask_class_id) if mask_class_id is not None else None,
        instances=instances,
    )


def load_yolo_dataset(images_dir, labels_dir, *, dataset_root: Path | None = None):
    images_dir = Path(images_dir)
    labels_dir = Path(labels_dir)
    dataset_root = Path(dataset_root) if dataset_root is not None else labels_dir.parent.parent
    images: list[Path] = []
    for ext in ("*.jpg", "*.jpeg", "*.png"):
        images.extend(images_dir.glob(ext))
    images = sorted(images)
    records = []
    for image_path in images:
        label_path = labels_dir / f"{image_path.stem}.txt"
        meta_path = labels_dir / f"{image_path.stem}.json"
        labels = []
        if label_path.exists():
            for line in label_path.read_text().splitlines():
                stripped = line.strip()
                if not stripped:
                    continue
                parts = stripped.split()
                if len(parts) < 5:
                    raise ValueError(f"invalid label line: {line}")
                class_id = int(float(parts[0]))
                coords = [float(value) for value in parts[1:5]]
                label: dict[str, Any] = {
                    "class_id": class_id,
                    "cx": coords[0],
                    "cy": coords[1],
                    "w": coords[2],
                    "h": coords[3],
                }

                # Optional: YOLO pose-style keypoints appended to the label line.
                # Supported:
                #   - 5 + 2*K: x1 y1 x2 y2 ...
                #   - 5 + 3*K: x1 y1 v1 x2 y2 v2 ...
                extra = parts[5:]
                if extra:
                    try:
                        extra_f = [float(v) for v in extra]
                    except Exception as exc:
                        raise ValueError(f"invalid label keypoints: {line}") from exc
                    label["keypoints"] = normalize_keypoints(extra_f, where="label.keypoints")

                labels.append(label)
        record = {"image": str(image_path), "labels": labels}
        record.update(_load_sidecar_metadata(meta_path, dataset_root))
        if not record["labels"]:
            record["labels"] = _derive_labels_from_mask(record)
        records.append(record)
    return records


def _pick_split(dataset_root: Path, split: str | None) -> str:
    if split:
        return split
    for candidate in ("val2017", "train2017"):
        if (dataset_root / "images" / candidate).exists():
            return candidate
    return "train2017"


def _resolve_dataset_json_layout(dataset_root: Path, split: str | None) -> tuple[Path, Path, str] | None:
    """Resolve images/labels directories from a dataset.json descriptor (optional).

    Some prepare tools write a dataset.json with absolute paths to avoid copying
    large image trees. When present, prefer it as a hint for build_manifest().
    """

    descriptor = dataset_root / "dataset.json"
    if not descriptor.exists():
        return None
    try:
        data = json.loads(descriptor.read_text(encoding="utf-8"))
    except Exception:
        return None
    if not isinstance(data, dict):
        return None

    images_dir = data.get("images_dir")
    labels_dir = data.get("labels_dir")
    split_from_desc = data.get("split")
    if not images_dir or not labels_dir:
        return None

    img = Path(images_dir)
    lbl = Path(labels_dir)
    if not img.is_absolute():
        img = dataset_root / img
    if not lbl.is_absolute():
        lbl = dataset_root / lbl

    effective_split = str(split_from_desc) if split_from_desc else None
    if split:
        effective_split = str(split)
        img_candidate = img.parent / effective_split
        lbl_candidate = lbl.parent / effective_split
        if img_candidate.exists():
            img = img_candidate
        if lbl_candidate.exists():
            lbl = lbl_candidate
    if effective_split is None:
        # Best-effort: infer from labels dir name.
        effective_split = lbl.name

    return img, lbl, effective_split


def build_manifest(dataset_root, *, split: str | None = None):
    dataset_root = Path(dataset_root)
    resolved = _resolve_dataset_json_layout(dataset_root, split)
    if resolved is not None:
        images_dir, labels_dir, split_effective = resolved
        records = load_yolo_dataset(images_dir, labels_dir, dataset_root=dataset_root)
        return {"images": records, "split": split_effective}

    split_effective = _pick_split(dataset_root, split)
    images_dir = dataset_root / "images" / split_effective
    labels_dir = dataset_root / "labels" / split_effective
    records = load_yolo_dataset(images_dir, labels_dir, dataset_root=dataset_root)
    return {"images": records, "split": split_effective}
