from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from ..write_data.writer import FieldValue, RESERVED_KEYS

EXTRA_TAGS_KEY = "extra_tags"
EXTRA_FIELDS_KEY = "extra_fields"
IGNORED_QOI_MAP_KEYS = frozenset({"fit_result", "fit_msg", "plot_msg"})
DEFAULT_KEEP_KEYS = frozenset(set(RESERVED_KEYS) | {EXTRA_TAGS_KEY, EXTRA_FIELDS_KEY})


def split_record(
    raw_record: Mapping[str, object],
) -> tuple[
    dict[str, object], Mapping[str, str] | None, Mapping[str, FieldValue] | None
]:
    """Split one raw record into base record, extra tags, and extra fields.

    Parameters
    ----------
    raw_record : Mapping[str, object]
        Input record containing canonical keys and optional ``extra_tags`` /
        ``extra_fields`` mappings.

    Returns
    -------
    tuple[dict[str, object], Mapping[str, str] | None, Mapping[str, FieldValue] | None]
        ``(record, extra_tags, extra_fields)`` tuple.

    Raises
    ------
    ValueError
        If ``extra_tags`` or ``extra_fields`` has an invalid shape.
    """
    record = _normalize_source_keys(raw_record)
    extra_tags = record.pop(EXTRA_TAGS_KEY, None)
    extra_fields = record.pop(EXTRA_FIELDS_KEY, None)
    return record, _as_string_map(extra_tags), _as_field_value_map(extra_fields)


def payload_to_records(payload: object) -> list[dict[str, object]]:
    """Normalize a JSON payload into a list of record dictionaries.

    Parameters
    ----------
    payload : object
        Parsed JSON payload. Supported forms are:
        - list of record mappings
        - single record mapping containing ``qoi``
        - QOI map where each key is a measurement name

    Returns
    -------
    list[dict[str, object]]
        Normalized record dictionaries.

    Raises
    ------
    ValueError
        If payload is not a mapping or list of mappings.
    """
    if isinstance(payload, list):
        records: list[dict[str, object]] = []
        for item in payload:
            if not isinstance(item, Mapping):
                continue
            records.append(_normalize_source_keys(item))
        return records

    if isinstance(payload, Mapping):
        payload_map = _normalize_source_keys(payload)
        if "qoi" in payload_map:
            return [payload_map]

        records: list[dict[str, object]] = []
        for qoi_name, value in payload_map.items():
            if qoi_name in IGNORED_QOI_MAP_KEYS:
                continue
            records.append(_record_from_qoi_map_value(str(qoi_name), value))
        return records

    raise ValueError("Unsupported JSON payload. Use dict or list of dicts.")


def _record_from_qoi_map_value(qoi_name: str, value: object) -> dict[str, object]:
    """Build one canonical record from a QOI-map entry.

    Parameters
    ----------
    qoi_name : str
        Measurement name.
    value : object
        Raw value from payload map.

    Returns
    -------
    dict[str, object]
        Canonical record dictionary containing at least ``qoi``.
    """
    if isinstance(value, Mapping):
        value_map = _normalize_source_keys(value)
        if value_map.get("__dtype__") == "UFloat":
            return {
                "qoi": qoi_name,
                "value": value_map.get("value", value_map.get("nominal_value")),
                "uncertainty": value_map.get("std_dev"),
            }

        record: dict[str, object] = {"qoi": qoi_name}
        for key in RESERVED_KEYS:
            if key == "qoi":
                continue
            if key in value_map:
                record[key] = value_map[key]
        if EXTRA_TAGS_KEY in value_map:
            record[EXTRA_TAGS_KEY] = value_map[EXTRA_TAGS_KEY]
        if EXTRA_FIELDS_KEY in value_map:
            record[EXTRA_FIELDS_KEY] = value_map[EXTRA_FIELDS_KEY]
        return record

    return {"qoi": qoi_name, "value": value}


def _normalize_source_keys(raw: Mapping[Any, Any]) -> dict[str, object]:
    """Convert arbitrary mapping keys to strings.

    Parameters
    ----------
    raw : Mapping[Any, Any]
        Source mapping.

    Returns
    -------
    dict[str, object]
        Mapping with string keys.
    """
    normalized: dict[str, object] = {}
    for key, value in raw.items():
        normalized[str(key)] = value
    return normalized


def _as_string_map(value: object) -> Mapping[str, str] | None:
    """Validate and normalize a string mapping.

    Parameters
    ----------
    value : object
        Candidate mapping.

    Returns
    -------
    Mapping[str, str] | None
        Normalized mapping, or ``None``.

    Raises
    ------
    ValueError
        If value is not a mapping.
    """
    if value is None:
        return None
    if not isinstance(value, Mapping):
        raise ValueError("'extra_tags' must be a mapping of strings.")
    return {str(key): str(item) for key, item in value.items()}


def _as_field_value_map(value: object) -> Mapping[str, FieldValue] | None:
    """Validate and normalize a mapping of supported field values.

    Parameters
    ----------
    value : object
        Candidate mapping.

    Returns
    -------
    Mapping[str, FieldValue] | None
        Normalized mapping, or ``None``.

    Raises
    ------
    ValueError
        If value is not a mapping or contains unsupported field types.
    """
    if value is None:
        return None
    if not isinstance(value, Mapping):
        raise ValueError("'extra_fields' must be a mapping.")

    normalized: dict[str, FieldValue] = {}
    for key, item in value.items():
        if not isinstance(item, (str, int, float, bool)):
            raise ValueError(
                f"Unsupported extra field type for key '{key}': {type(item)}"
            )
        normalized[str(key)] = item
    return normalized
