from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

from .utils import (
    future_offsets,
    load_yaml,
    parse_time,
    read_json,
    source_existing_dat_path,
    source_meta_path,
    status_columns,
    time_offsets,
)


class SourceStore:
    def __init__(self, output_root: str | Path, source: str):
        self.output_root = Path(output_root)
        self.source = source
        self.dat_path = source_existing_dat_path(self.output_root, source)
        self.meta_path = source_meta_path(self.dat_path)
        self.meta = read_json(self.meta_path)
        self.dtype = np.dtype(self.meta["dtype"])
        self.row_shape = tuple(int(v) for v in self.meta["row_shape"])
        self.row_count = int(self.meta["row_count"])
        self.shape = (self.row_count, *self.row_shape)
        self._mm = np.memmap(self.dat_path, dtype=self.dtype, mode="r", shape=self.shape)

    def get(self, idx: int) -> np.ndarray:
        idx = int(idx)
        if idx < 0 or idx >= self.row_count:
            raise IndexError(f"{self.source} index out of range: {idx} / {self.row_count}")
        return np.asarray(self._mm[idx], dtype=np.float32)


class FastDataset(Dataset):
    def __init__(
        self,
        config: str | Path | dict[str, Any],
        split: str | None = None,
        required_inputs: list[str] | None = None,
        required_labels: list[str] | None = None,
        dtype: torch.dtype = torch.float32,
    ):
        self.config = load_yaml(config) if not isinstance(config, dict) else dict(config)
        self.output_root = Path(self.config["output_root"])
        self.catalog_path = Path(self.config.get("catalog_path", self.output_root / "catalog.csv"))
        self.catalog = pd.read_csv(self.catalog_path, dtype={"timestamp": str})
        self.catalog["timestamp"] = self.catalog["timestamp"].astype(str)

        self.required_inputs = list(required_inputs if required_inputs is not None else self.config.get("required_inputs", ["concat"]))
        self.required_labels = list(required_labels if required_labels is not None else self.config.get("required_labels", ["ci"]))
        self.dtype = dtype

        window = self.config.get("input_window", {})
        self.input_offsets = list(window.get("offset_minutes") or time_offsets(
            int(window.get("past_minutes", 50)),
            int(window.get("interval_minutes", 10)),
        ))
        bt_cfg = self.config.get("labels", {}).get("bt", {})
        self.bt_source = str(bt_cfg.get("source", bt_cfg.get("source_name", "bt")))
        self.use_bt_mask = bool(bt_cfg.get("use_mask", True))
        self.bt_mask_source = str(bt_cfg.get("mask_source", "bt_mask"))
        self.bt_mask_regions = tuple(str(region).lower() for region in bt_cfg.get("mask_regions", ["object", "context"]))
        valid_mask_regions = {"object", "context"}
        unknown_mask_regions = sorted(set(self.bt_mask_regions) - valid_mask_regions)
        if unknown_mask_regions:
            raise ValueError(f"unsupported BT mask regions: {unknown_mask_regions}")
        if self.use_bt_mask and not self.bt_mask_regions:
            raise ValueError("labels.bt.mask_regions must contain at least one of: object, context")
        self.bt_offsets = list(bt_cfg.get("lead_minutes") or future_offsets(
            int(bt_cfg.get("future_minutes", 60)),
            int(bt_cfg.get("interval_minutes", 10)),
        ))

        self.time_to_row = {str(row.timestamp): i for i, row in self.catalog.iterrows()}
        store_sources = {self._catalog_source_name(source) for source in self.required_inputs + self.required_labels}
        if "bt" in self.required_labels and self.use_bt_mask:
            store_sources.add(self.bt_mask_source)
        available_sources = self._available_sources()
        self.stores = {
            source: SourceStore(self.output_root, source)
            for source in store_sources
            if source in available_sources
        }
        self.samples = self._build_samples(split)

    def _catalog_source_name(self, name: str) -> str:
        if name == "ci":
            return "ci_hard"
        if name == "bt":
            return self.bt_source
        return name

    def _available_sources(self) -> set[str]:
        out = set()
        for col in self.catalog.columns:
            if col.endswith("_status"):
                out.add(col[: -len("_status")])
        return out

    def _filter_split(self, times: list[str], split: str | None) -> list[str]:
        if not split:
            return times
        ranges = self.config.get("splits", {}).get(split)
        if ranges is None:
            raise KeyError(f"split {split!r} not found in config.splits")
        parsed_ranges = []
        for item in ranges:
            start, end = (item["start"], item["end"]) if isinstance(item, dict) else item
            parsed_ranges.append((parse_time(start), parse_time(end)))
        kept = []
        for ts in times:
            t = parse_time(ts)
            if any(start <= t <= end for start, end in parsed_ranges):
                kept.append(ts)
        return kept

    def _is_ok(self, timestamp: str, source: str) -> bool:
        row_idx = self.time_to_row.get(timestamp)
        if row_idx is None:
            return False
        idx_col, status_col = status_columns(source)
        row = self.catalog.iloc[row_idx]
        return status_col in row and row[status_col] == "ok" and not pd.isna(row[idx_col])

    def _idx_for(self, timestamp: str, source: str) -> int:
        row_idx = self.time_to_row[timestamp]
        idx_col, status_col = status_columns(source)
        row = self.catalog.iloc[row_idx]
        if row[status_col] != "ok" or pd.isna(row[idx_col]):
            raise RuntimeError(f"catalog says source not ok: {timestamp} {source} status={row[status_col]}")
        return int(row[idx_col])

    def _input_times(self, sample_time: str) -> list[str]:
        base = parse_time(sample_time)
        return [(base + np.timedelta64(int(offset), "m")).strftime("%Y%m%d%H%M") for offset in self.input_offsets]

    def _bt_times(self, sample_time: str) -> list[str]:
        base = parse_time(sample_time)
        return [(base + np.timedelta64(int(offset), "m")).strftime("%Y%m%d%H%M") for offset in self.bt_offsets]

    def _valid_sample(self, sample_time: str) -> bool:
        for source in self.required_inputs:
            catalog_source = self._catalog_source_name(source)
            for ts in self._input_times(sample_time):
                if not self._is_ok(ts, catalog_source):
                    return False
        for label in self.required_labels:
            catalog_source = self._catalog_source_name(label)
            if label == "bt":
                if self.use_bt_mask and not self._is_ok(sample_time, self.bt_mask_source):
                    return False
                for ts in self._bt_times(sample_time):
                    if not self._is_ok(ts, catalog_source):
                        return False
            else:
                if not self._is_ok(sample_time, catalog_source):
                    return False
        return True

    def _build_samples(self, split: str | None) -> list[str]:
        times = sorted(self.catalog["timestamp"].astype(str).tolist())
        times = self._filter_split(times, split)
        return [ts for ts in times if self._valid_sample(ts)]

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

    def _load_input(self, source: str, sample_time: str) -> torch.Tensor:
        catalog_source = self._catalog_source_name(source)
        store = self.stores[catalog_source]
        rows = [store.get(self._idx_for(ts, catalog_source)) for ts in self._input_times(sample_time)]
        return torch.from_numpy(np.stack(rows, axis=0).astype(np.float32, copy=False)).to(dtype=self.dtype)

    def _load_label(self, label: str, sample_time: str) -> torch.Tensor:
        catalog_source = self._catalog_source_name(label)
        store = self.stores[catalog_source]
        if label == "bt":
            rows = [store.get(self._idx_for(ts, catalog_source)) for ts in self._bt_times(sample_time)]
            bt = np.stack(rows, axis=0).astype(np.float32, copy=False)
            if self.use_bt_mask:
                packed_mask = self.stores[self.bt_mask_source].get(
                    self._idx_for(sample_time, self.bt_mask_source)
                )
                bt = self.apply_bt_mask(bt, packed_mask)
            return torch.from_numpy(bt.astype(np.float32, copy=False)).to(dtype=self.dtype)
        if label == "bt_mask":
            arr = store.get(self._idx_for(sample_time, catalog_source))
            return torch.from_numpy(arr.astype(np.float32, copy=False)).to(dtype=self.dtype)
        arr = store.get(self._idx_for(sample_time, catalog_source))
        return torch.from_numpy(arr.astype(np.float32, copy=False)).to(dtype=self.dtype)

    def decode_bt_mask(self, mask_row: np.ndarray) -> np.ndarray:
        """Return one boolean BT-loss mask per configured forecast lead."""
        if self.bt_mask_source == "bt_mask":
            static_mask = np.asarray(mask_row, dtype=bool)
            if static_mask.ndim != 2:
                raise ValueError(f"bt_mask row must be 2-D, got {static_mask.shape}")
            return np.repeat(static_mask[np.newaxis, ...], len(self.bt_offsets), axis=0)

        if self.bt_mask_source != "bt_lineage_mask":
            raise ValueError(f"unsupported labels.bt.mask_source: {self.bt_mask_source!r}")

        packed = np.asarray(mask_row, dtype=np.uint16)
        if packed.ndim != 2:
            raise ValueError(f"bt_lineage_mask row must be 2-D, got {packed.shape}")

        meta = self.stores[self.bt_mask_source].meta
        object_bits = {int(lead): int(bit) for lead, bit in dict(meta.get("object_bits", {})).items()}
        context_bits = {int(lead): int(bit) for lead, bit in dict(meta.get("context_bits", {})).items()}
        missing_leads = [int(lead) for lead in self.bt_offsets if int(lead) not in object_bits or int(lead) not in context_bits]
        if missing_leads:
            raise ValueError(
                f"bt_lineage_mask metadata has no object/context bits for leads {missing_leads}"
            )

        lead_masks = []
        for lead in self.bt_offsets:
            valid = np.zeros(packed.shape, dtype=bool)
            if "object" in self.bt_mask_regions:
                valid |= (packed & np.uint16(1 << object_bits[int(lead)])) != 0
            if "context" in self.bt_mask_regions:
                valid |= (packed & np.uint16(1 << context_bits[int(lead)])) != 0
            lead_masks.append(valid)
        return np.stack(lead_masks, axis=0)

    def apply_bt_mask(self, bt: np.ndarray, mask_row: np.ndarray) -> np.ndarray:
        """Replace pixels outside the configured lead-specific BT mask with NaN."""
        bt = np.asarray(bt, dtype=np.float32)
        valid = self.decode_bt_mask(mask_row)
        expected_shape = (bt.shape[0], *bt.shape[-2:])
        if tuple(valid.shape) != expected_shape:
            raise ValueError(f"BT/mask shape mismatch: bt={bt.shape}, mask={valid.shape}")
        return np.where(valid[:, np.newaxis, ...], bt, np.nan)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        sample_time = self.samples[int(idx)]
        return {
            "inputs": {name: self._load_input(name, sample_time) for name in self.required_inputs},
            "labels": {name: self._load_label(name, sample_time) for name in self.required_labels},
            "time": sample_time,
        }