from __future__ import annotations

import hashlib
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator, Mapping, Sequence

import numpy as np

from pyimgano.utils.path_normalize import normalize_path


@dataclass(frozen=True)
class ManifestSplitPolicy:
    """Controls how missing `split` fields are auto-assigned.

    Notes
    -----
    - v1 is intentionally narrow: industrial benchmarks usually assume that
      anomalies belong in test, while normals are split between train/test.
    - Splits are deterministic given the same manifest + seed.
    """

    mode: str = "benchmark"
    scope: str = "category"
    seed: int = 0
    test_normal_fraction: float = 0.2


@dataclass(frozen=True)
class ManifestRecord:
    image_path: str
    category: str
    split: str | None = None
    label: int | None = None
    mask_path: str | None = None
    group_id: str | None = None
    meta: Mapping[str, Any] | None = None

    @classmethod
    def from_mapping(cls, raw: Mapping[str, Any], *, lineno: int) -> "ManifestRecord":
        def _require_str(key: str) -> str:
            value = raw.get(key, None)
            if value is None:
                raise ValueError(f"Manifest line {lineno}: missing required field {key!r}.")
            s = str(value).strip()
            if not s:
                raise ValueError(f"Manifest line {lineno}: field {key!r} must be non-empty.")
            return s

        image_path = _require_str("image_path")
        category = _require_str("category")

        split_raw = raw.get("split", None)
        split = None if split_raw is None else str(split_raw).strip().lower()
        if split is not None and split not in ("train", "val", "test"):
            raise ValueError(
                f"Manifest line {lineno}: invalid split {split!r}. "
                "Supported: train, val, test."
            )

        label_raw = raw.get("label", None)
        label = None
        if label_raw is not None:
            try:
                label_int = int(label_raw)
            except Exception as exc:  # noqa: BLE001 - validation boundary
                raise ValueError(
                    f"Manifest line {lineno}: label must be 0/1, got {label_raw!r}."
                ) from exc
            if label_int not in (0, 1):
                raise ValueError(f"Manifest line {lineno}: label must be 0/1, got {label_int!r}.")
            label = label_int

        mask_path_raw = raw.get("mask_path", None)
        mask_path = None if mask_path_raw is None else str(mask_path_raw).strip()
        if mask_path is not None and not mask_path:
            mask_path = None

        group_id_raw = raw.get("group_id", None)
        group_id = None if group_id_raw is None else str(group_id_raw).strip()
        if group_id is not None and not group_id:
            group_id = None

        meta_raw = raw.get("meta", None)
        meta: Mapping[str, Any] | None
        if meta_raw is None:
            meta = None
        elif isinstance(meta_raw, Mapping):
            meta = meta_raw
        else:
            raise ValueError(
                f"Manifest line {lineno}: meta must be an object/dict, got {type(meta_raw).__name__}."
            )

        # Explicit-split validation rules.
        if split == "test" and label is None:
            raise ValueError(f"Manifest line {lineno}: split='test' requires an explicit label.")
        if split in ("train", "val") and label == 1:
            raise ValueError(
                f"Manifest line {lineno}: split={split!r} cannot have label=1 (anomaly)."
            )

        return cls(
            image_path=image_path,
            category=category,
            split=split,
            label=label,
            mask_path=mask_path,
            group_id=group_id,
            meta=meta,
        )


def iter_manifest_records(manifest_path: str | Path) -> Iterator[ManifestRecord]:
    path = Path(manifest_path)
    if not path.exists():
        raise FileNotFoundError(f"Manifest not found: {path}")

    with path.open("r", encoding="utf-8") as f:
        for i, line in enumerate(f, start=1):
            text = line.strip()
            if not text:
                continue
            if text.startswith("#"):
                continue
            try:
                raw = json.loads(text)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Manifest line {i}: invalid JSON ({exc.msg}).") from exc
            if not isinstance(raw, Mapping):
                raise ValueError(
                    f"Manifest line {i}: expected a JSON object, got {type(raw).__name__}."
                )
            yield ManifestRecord.from_mapping(raw, lineno=i)


def list_manifest_categories(manifest_path: str | Path) -> list[str]:
    cats: set[str] = set()
    for rec in iter_manifest_records(manifest_path):
        cats.add(str(rec.category))
    return sorted(cats)


def _seed_for_category(seed: int, category: str) -> int:
    payload = f"{int(seed)}:{category}".encode("utf-8")
    digest = hashlib.sha256(payload).digest()
    return int.from_bytes(digest[:8], "little", signed=False) % (2**63 - 1)


def _resolve_existing_path(
    raw_value: str,
    *,
    manifest_path: Path,
    root_fallback: Path | None,
) -> Path:
    raw = normalize_path(str(raw_value))
    p = Path(raw)
    if p.is_absolute():
        if p.exists():
            return p
        raise FileNotFoundError(f"Path not found: {p}")

    cand1 = (manifest_path.parent / p).resolve()
    if cand1.exists():
        return cand1

    if root_fallback is not None:
        cand2 = (root_fallback / p).resolve()
        if cand2.exists():
            return cand2

    raise FileNotFoundError(
        "Path not found. Tried: "
        f"{cand1} (manifest-dir) and {((root_fallback / p).resolve() if root_fallback else '<no root>')} (root)."
    )


@dataclass(frozen=True)
class ManifestBenchmarkSplit:
    train_paths: list[str]
    calibration_paths: list[str]
    test_paths: list[str]
    test_labels: np.ndarray
    test_masks: np.ndarray | None
    pixel_skip_reason: str | None = None
    test_meta: list[Mapping[str, Any] | None] | None = None


def load_manifest_benchmark_split(
    *,
    manifest_path: str | Path,
    root_fallback: str | Path | None,
    category: str,
    resize: tuple[int, int] = (256, 256),
    load_masks: bool | str = True,
    missing_mask_policy: str = "skip",
    split_policy: ManifestSplitPolicy | None = None,
) -> ManifestBenchmarkSplit:
    """Load a benchmark-style split from a JSONL manifest.

    Parameters
    ----------
    manifest_path:
        JSONL file path.
    root_fallback:
        Used only for resolving relative `image_path`/`mask_path` values when the
        file does not exist relative to the manifest directory.
    category:
        Category name.
    resize:
        Target (H, W) for masks (and for consistency with benchmark loaders).
    load_masks:
        - When false, masks are not loaded and pixel metrics can’t be computed.
        - When true, masks are loaded when available.
        - When "auto", masks are loaded only when at least one `mask_path` is present
          in the resolved test split; otherwise pixel metrics are disabled.
    missing_mask_policy:
        Policy when `load_masks` is enabled but anomaly masks are missing.
        - "skip" (default): disable pixel metrics and return `pixel_skip_reason`.
        - "error": raise a ValueError.
    split_policy:
        Auto split policy used when records are missing `split`.
    """

    mp = Path(manifest_path)
    root_path = None if root_fallback is None else Path(root_fallback)
    policy = split_policy or ManifestSplitPolicy()

    missing_mask_policy_norm = str(missing_mask_policy).strip().lower()
    if missing_mask_policy_norm not in ("skip", "error"):
        raise ValueError(
            f"missing_mask_policy must be one of: skip, error. Got: {missing_mask_policy!r}"
        )

    load_masks_mode = load_masks
    load_masks_norm: str | None = None
    if isinstance(load_masks_mode, str):
        load_masks_norm = str(load_masks_mode).strip().lower()
        if load_masks_norm != "auto":
            raise ValueError("load_masks must be a bool or 'auto'.")

    cat = str(category)
    records: list[ManifestRecord] = [r for r in iter_manifest_records(mp) if str(r.category) == cat]
    if not records:
        raise ValueError(f"Manifest contains no records for category={cat!r}.")

    if policy.scope != "category":
        raise ValueError(
            f"Only split_policy.scope='category' is supported for v1. Got: {policy.scope!r}."
        )
    if policy.mode != "benchmark":
        raise ValueError(
            f"Only split_policy.mode='benchmark' is supported for v1. Got: {policy.mode!r}."
        )
    frac = float(policy.test_normal_fraction)
    if not (0.0 <= frac <= 1.0):
        raise ValueError(
            f"split_policy.test_normal_fraction must be in [0,1]. Got: {policy.test_normal_fraction!r}."
        )

    def _group_id_for_record(rec: ManifestRecord, *, idx: int) -> str:
        gid = rec.group_id
        meta_gid = None
        if rec.meta is not None:
            try:
                meta_gid = rec.meta.get("group_id", None)
            except Exception:  # noqa: BLE001 - best-effort metadata probing
                meta_gid = None

        if gid is not None and meta_gid is not None:
            a = str(gid).strip()
            b = str(meta_gid).strip()
            if a and b and a != b:
                raise ValueError(
                    "Manifest record has conflicting group identifiers: "
                    f"group_id={a!r} meta.group_id={b!r}."
                )

        if gid is not None and str(gid).strip():
            return str(gid).strip()
        if meta_gid is not None and str(meta_gid).strip():
            return str(meta_gid).strip()
        return f"__ungrouped__{idx}"

    # Group records (group-aware split when group_id/meta.group_id present).
    group_to_indices: dict[str, list[int]] = {}
    for idx, rec in enumerate(records):
        g = _group_id_for_record(rec, idx=idx)
        group_to_indices.setdefault(str(g), []).append(int(idx))

    group_ids = list(group_to_indices.keys())

    # Determine fixed group split, and list normal candidate groups.
    fixed_train: set[str] = set()
    fixed_val: set[str] = set()
    fixed_test: set[str] = set()
    normal_candidate: list[str] = []

    for gid in group_ids:
        idxs = group_to_indices[gid]
        splits = {records[i].split for i in idxs if records[i].split is not None}
        if len(splits) > 1:
            raise ValueError(
                f"Manifest group_id={gid!r} has conflicting explicit splits: {sorted(splits)}."
            )
        explicit_split = next(iter(splits), None)

        # Any anomaly label in the group forces test.
        has_anomaly = any(records[i].label == 1 for i in idxs if records[i].label is not None)
        if explicit_split in ("train", "val") and has_anomaly:
            raise ValueError(
                f"Manifest group_id={gid!r} has split={explicit_split!r} but contains label=1 records."
            )

        if explicit_split == "train":
            fixed_train.add(gid)
            continue
        if explicit_split == "val":
            fixed_val.add(gid)
            continue
        if explicit_split == "test":
            fixed_test.add(gid)
            continue

        # No explicit split: apply benchmark policy.
        if has_anomaly:
            fixed_test.add(gid)
        else:
            normal_candidate.append(gid)

    # Select a fraction of normal candidate groups into test.
    normal_candidate_counts = [len(group_to_indices[gid]) for gid in normal_candidate]
    total_normals = int(sum(normal_candidate_counts))
    target_test_normals = int(math.ceil(total_normals * frac))
    if 0.0 < frac < 1.0 and total_normals >= 2:
        target_test_normals = max(1, min(total_normals - 1, target_test_normals))
    target_test_normals = max(0, min(total_normals, target_test_normals))
    # If the split would produce an empty train set (no explicit train items),
    # bias toward keeping at least one normal sample for training.
    if total_normals > 0 and not fixed_train and target_test_normals >= total_normals:
        target_test_normals = max(0, total_normals - 1)

    selected_test_normal: set[str] = set()
    if target_test_normals > 0 and normal_candidate:
        rng = np.random.default_rng(_seed_for_category(int(policy.seed), cat))
        order = list(range(len(normal_candidate)))
        rng.shuffle(order)
        running = 0
        for j in order:
            gid = normal_candidate[j]
            if gid in selected_test_normal:
                continue
            selected_test_normal.add(gid)
            running += int(len(group_to_indices[gid]))
            if running >= target_test_normals:
                break

    # Assemble outputs preserving original record order.
    train_paths: list[str] = []
    cal_paths: list[str] = []
    test_paths: list[str] = []
    test_labels: list[int] = []
    test_mask_paths: list[str | None] = []
    test_meta: list[Mapping[str, Any] | None] = []
    missing_anomaly_mask = False
    any_mask_field_present = False

    for idx, rec in enumerate(records):
        gid = str(_group_id_for_record(rec, idx=idx))

        if gid in fixed_val:
            cal_paths.append(
                str(
                    _resolve_existing_path(
                        rec.image_path, manifest_path=mp, root_fallback=root_path
                    )
                )
            )
            continue

        in_test = gid in fixed_test or gid in selected_test_normal
        if not in_test and gid in fixed_train:
            train_paths.append(
                str(
                    _resolve_existing_path(
                        rec.image_path, manifest_path=mp, root_fallback=root_path
                    )
                )
            )
            continue
        if not in_test and gid not in fixed_test and gid not in selected_test_normal:
            # Unspecified normal group: default to train.
            train_paths.append(
                str(
                    _resolve_existing_path(
                        rec.image_path, manifest_path=mp, root_fallback=root_path
                    )
                )
            )
            continue

        # Test record (explicit or assigned).
        label = int(rec.label) if rec.label is not None else 0
        test_paths.append(
            str(_resolve_existing_path(rec.image_path, manifest_path=mp, root_fallback=root_path))
        )
        test_labels.append(label)
        if rec.meta is None:
            test_meta.append(None)
        else:
            test_meta.append(dict(rec.meta))

        if load_masks_norm is None and not bool(load_masks_mode):
            test_mask_paths.append(None)
            continue

        if rec.mask_path is not None and str(rec.mask_path).strip():
            any_mask_field_present = True

        if label == 1:
            if rec.mask_path is None:
                missing_anomaly_mask = True
                test_mask_paths.append(None)
            else:
                try:
                    resolved = _resolve_existing_path(
                        rec.mask_path, manifest_path=mp, root_fallback=root_path
                    )
                except FileNotFoundError:
                    missing_anomaly_mask = True
                    test_mask_paths.append(None)
                else:
                    test_mask_paths.append(str(resolved))
        else:
            # Normal: masks optional; treat missing as background.
            if rec.mask_path is None:
                test_mask_paths.append(None)
            else:
                try:
                    resolved = _resolve_existing_path(
                        rec.mask_path, manifest_path=mp, root_fallback=root_path
                    )
                except FileNotFoundError:
                    test_mask_paths.append(None)
                else:
                    test_mask_paths.append(str(resolved))

    if not train_paths:
        raise ValueError(f"Manifest category={cat!r} produced an empty train split.")
    if not test_paths:
        raise ValueError(f"Manifest category={cat!r} produced an empty test split.")

    masks_arr: np.ndarray | None = None
    pixel_skip_reason: str | None = None
    meta_out: list[Mapping[str, Any] | None] | None = None
    # Determine whether masks should be loaded.
    should_load_masks = bool(load_masks_mode) if load_masks_norm is None else True
    if load_masks_norm == "auto":
        should_load_masks = any_mask_field_present

    if not should_load_masks and load_masks_norm == "auto":
        pixel_skip_reason = "load_masks=auto skipped mask loading (no mask_path entries present)."

    if should_load_masks:
        has_any_masks = any(p is not None for p in test_mask_paths)

        if has_any_masks and any(
            lab == 1 for lab, p in zip(test_labels, test_mask_paths) if p is None and lab == 1
        ):
            missing_anomaly_mask = True

        if missing_anomaly_mask:
            msg = "Missing mask_path (or missing mask files) for anomaly test samples."
            if missing_mask_policy_norm == "error":
                raise ValueError(msg)
            pixel_skip_reason = msg
            masks_arr = None
        elif not has_any_masks:
            if any_mask_field_present:
                pixel_skip_reason = (
                    "mask_path entries present but no mask files were found on disk."
                )
            else:
                pixel_skip_reason = "No mask_path entries found in manifest test records."
            masks_arr = None
        else:
            masks_arr = _load_masks_or_zeros(test_mask_paths, resize=resize)

    if any(m is not None for m in test_meta):
        meta_out = test_meta

    return ManifestBenchmarkSplit(
        train_paths=train_paths,
        calibration_paths=cal_paths,
        test_paths=test_paths,
        test_labels=np.asarray(test_labels, dtype=np.int64),
        test_masks=masks_arr,
        pixel_skip_reason=pixel_skip_reason,
        test_meta=meta_out,
    )


def _load_masks_or_zeros(
    mask_paths: Sequence[str | None], *, resize: tuple[int, int]
) -> np.ndarray:
    import cv2

    h, w = int(resize[0]), int(resize[1])
    out: list[np.ndarray] = []
    for p in mask_paths:
        if p is None:
            out.append(np.zeros((h, w), dtype=np.uint8))
            continue
        arr = cv2.imread(str(p), cv2.IMREAD_GRAYSCALE)
        if arr is None:
            out.append(np.zeros((h, w), dtype=np.uint8))
            continue
        if arr.shape != (h, w):
            arr = cv2.resize(arr, (w, h), interpolation=cv2.INTER_NEAREST)
        out.append((arr > 127).astype(np.uint8))
    return np.stack(out, axis=0)
