from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterable, Iterator, Mapping

from ..contract_injection import ExperimentRowMapper, map_experiment_row
from ..contracts import ExperimentRow, stable_json_fingerprint

_RESERVED_FIELDS = frozenset(
    {
        "sample_id",
        "id",
        "payload",
        "ground_truth",
        "meta",
        "metadata",
        "quality",
    }
)


@dataclass(frozen=True)
class SageDataMappingContext:
    dataset_ref: str
    usage: str | None = None
    profile: str | None = None
    source_version: str | None = None
    snapshot_hash: str = ""
    source_format: str = "python"
    quality: dict[str, Any] = field(default_factory=dict)
    sample_id_field: str = "sample_id"
    fallback_id_field: str = "id"



def _resolve_sample_id(raw: Mapping[str, Any], context: SageDataMappingContext) -> str:
    sample_id = raw.get(context.sample_id_field)
    if isinstance(sample_id, str) and sample_id.strip():
        return sample_id.strip()

    fallback_id = raw.get(context.fallback_id_field)
    if isinstance(fallback_id, str) and fallback_id.strip():
        return fallback_id.strip()

    return stable_json_fingerprint(dict(raw))



def _resolve_payload(raw: Mapping[str, Any]) -> dict[str, Any]:
    direct_payload = raw.get("payload")
    if isinstance(direct_payload, Mapping):
        return dict(direct_payload)

    return {k: v for k, v in raw.items() if k not in _RESERVED_FIELDS}



def _resolve_ground_truth(raw: Mapping[str, Any]) -> dict[str, Any] | None:
    ground_truth = raw.get("ground_truth")
    if ground_truth is None:
        return None
    if isinstance(ground_truth, Mapping):
        return dict(ground_truth)
    return {"value": ground_truth}



def _resolve_quality(raw: Mapping[str, Any], context: SageDataMappingContext) -> dict[str, Any]:
    quality = dict(context.quality)

    raw_quality = raw.get("quality")
    if isinstance(raw_quality, Mapping):
        quality.update(dict(raw_quality))

    metadata = raw.get("metadata")
    if isinstance(metadata, Mapping):
        nested_quality = metadata.get("quality")
        if isinstance(nested_quality, Mapping):
            quality.update(dict(nested_quality))

    return quality



def default_sagedata_row_mapper(
    raw_row: Mapping[str, Any],
    *,
    context: SageDataMappingContext,
) -> dict[str, Any]:
    """Map one sageData-like row into ExperimentRow-compatible dict."""
    if not isinstance(raw_row, Mapping):
        raise TypeError(f"raw_row must be mapping-like, got {type(raw_row)}")

    row = dict(raw_row)
    return {
        "sample_id": _resolve_sample_id(row, context),
        "payload": _resolve_payload(row),
        "ground_truth": _resolve_ground_truth(row),
        "meta": {
            "dataset_ref": context.dataset_ref,
            "usage": context.usage,
            "profile": context.profile,
            "source_version": context.source_version,
            "snapshot_hash": context.snapshot_hash,
            "source_format": context.source_format,
            "quality": _resolve_quality(row, context),
        },
    }



def build_sagedata_mapper(
    *,
    context: SageDataMappingContext,
) -> ExperimentRowMapper:
    """Build a mapper closure for map_experiment_row(..., mapper=...)."""

    def _mapper(raw_row: Mapping[str, Any]) -> Mapping[str, Any]:
        return default_sagedata_row_mapper(raw_row, context=context)

    return _mapper



def map_sagedata_row(
    raw_row: Mapping[str, Any],
    *,
    context: SageDataMappingContext,
    mapper: ExperimentRowMapper | None = None,
) -> ExperimentRow:
    """Map one row into validated ExperimentRow using default or custom mapper."""
    if mapper is None:
        mapper = build_sagedata_mapper(context=context)
    return map_experiment_row(raw_row, mapper=mapper)



def iter_sagedata_rows(
    rows: Iterable[Mapping[str, Any]],
    *,
    context: SageDataMappingContext,
    mapper: ExperimentRowMapper | None = None,
) -> Iterator[ExperimentRow]:
    """Map iterable rows into validated ExperimentRow stream."""
    for raw_row in rows:
        yield map_sagedata_row(raw_row, context=context, mapper=mapper)


__all__ = [
    "SageDataMappingContext",
    "default_sagedata_row_mapper",
    "build_sagedata_mapper",
    "map_sagedata_row",
    "iter_sagedata_rows",
]
