from __future__ import annotations

from typing import Any, Mapping

from pydantic import BaseModel, ConfigDict, Field, PositiveInt

# Reuse canonical definitions from the video schema module
from plume_nav_sim.video.schema import DIMS_TYX, SCHEMA_VERSION
from plume_nav_sim.video.schema import VARIABLE_NAME as _VIDEO_VARIABLE_NAME
from plume_nav_sim.video.schema import VideoPlumeAttrs
from plume_nav_sim.video.schema import validate_attrs as _validate_video_attrs

# Backwards-compatible export using the canonical variable name
VARIABLE_CONCENTRATION: str = _VIDEO_VARIABLE_NAME

# Optional Timebase remains defined here for loaders that use it; attrs model
# unification is handled via plume_nav_sim.video.schema.VideoPlumeAttrs.


class Timebase(BaseModel):
    """Rational timebase for frame timestamps (numerator/denominator in Hz)."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    numerator: PositiveInt
    denominator: PositiveInt = Field(default=1)


def validate_attrs(attrs: Mapping[str, Any]) -> VideoPlumeAttrs:
    # Delegate to the canonical validator from the video schema module so callers
    # receive the canonical VideoPlumeAttrs type.
    return _validate_video_attrs(attrs)


def validate_xarray_like(ds: Any) -> VideoPlumeAttrs:
    # Resolve data variable by name from either mapping or attribute access
    data_vars = getattr(ds, "data_vars", None)
    var = None
    if isinstance(data_vars, Mapping):
        var = data_vars.get(VARIABLE_CONCENTRATION)
    if var is None and hasattr(ds, "__getitem__"):
        try:
            var = ds[VARIABLE_CONCENTRATION]
        except Exception:
            var = None
    if var is None:
        raise ValueError(f"Missing required data variable '{VARIABLE_CONCENTRATION}'")

    # Validate dims. Prefer explicit array attribute if present; otherwise use
    # the variable's dims. Enforce exact match with DIMS_TYX.
    declared_dims = None
    var_attrs = getattr(var, "attrs", None)
    if isinstance(var_attrs, Mapping) and "_ARRAY_DIMENSIONS" in var_attrs:
        declared_dims = var_attrs.get("_ARRAY_DIMENSIONS")
    if declared_dims is None:
        declared_dims = getattr(var, "dims", None)

    if not isinstance(declared_dims, (tuple, list)) or tuple(declared_dims) != DIMS_TYX:
        raise ValueError(
            f"Variable '{VARIABLE_CONCENTRATION}' must have dims {DIMS_TYX}, got {declared_dims!r}"
        )

    attrs = getattr(ds, "attrs", None)
    if not isinstance(attrs, Mapping):
        raise ValueError("Dataset attrs must be a mapping of metadata")

    merged_attrs: dict[str, Any] = dict(attrs)
    if isinstance(var_attrs, Mapping):
        merged_attrs = {**dict(var_attrs), **merged_attrs}

    return validate_attrs(merged_attrs)


__all__ = [
    "VARIABLE_CONCENTRATION",
    "DIMS_TYX",
    "SCHEMA_VERSION",
    "Timebase",
    "VideoPlumeAttrs",
    "validate_attrs",
    "validate_xarray_like",
]
