from __future__ import annotations

from collections.abc import Mapping, Sequence
import math
from typing import Any

from ..write_data.writer import FieldValue, RESERVED_KEYS

EXTRA_TAGS_KEY = "extra_tags"
EXTRA_FIELDS_KEY = "extra_fields"
SNAPSHOT_INSTRUMENTS_KEY = "instruments"
SNAPSHOT_SUBMODULES_KEY = "submodules"
SNAPSHOT_PARAMETERS_KEY = "parameters"
SNAPSHOT_QOI_MODULE_KEY = "qoi"
SNAPSHOT_QUANTUM_DEVICE_KEY = "quantum_device"
IGNORED_QOI_MAP_KEYS = frozenset(
    {
        "best_fit",
        "fit_success",
        "fit_result",
        "fit_msg",
        "plot_msg",
    }
)
PARTNER_EXTRA_TAG_ALIASES = {
    "area_of_interest": "area_of_interest",
    "long_label": "long_label",
}
PARTNER_EXTRA_TAG_KEYS = frozenset(PARTNER_EXTRA_TAG_ALIASES)
DEFAULT_KEEP_KEYS = frozenset(
    set(RESERVED_KEYS)
    | {EXTRA_TAGS_KEY, EXTRA_FIELDS_KEY}
    | set(PARTNER_EXTRA_TAG_ALIASES.values())
    | PARTNER_EXTRA_TAG_KEYS
)
SCQT_TRANSMON_QOI_KEYS = frozenset(
    {
        "resonator_qc",
        "resonator_qi",
        "resonator_freq_low",
        "resonator_freq_high",
        "qubit_01_frequency",
        "qubit_12_frequency",
        "anharmonicity",
        "pi_pulse_duration",
        "pi_pulse_amp",
        "motzoi",
        "readout_amplitude",
        "readout_pulse_duration",
        "flux_quantum_current",
        "ac_flux_periodicity",
        "t1",
        "t2",
        "t2_echo",
        "single_qubit_fidelity_simultaneous",
        "single_qubit_fidelity_isolated",
        "ssro_fidelity",
        "dispersive_shift_power",
        "dispersive_shift_state",
        "sweetspot_current_dc",
        "charging_energy",
        "josephson_energy",
        "qubit_coupling",
        "drive_sensitivity",
        "flux_sensitivity",
        "effective_temperature",
        "flux_noise_scaling_factor_ramsey",
        "flux_noise_scaling_factor_echo",
        "residual_photon_number",
    }
)
QOI_ANALYSIS_ALIASES = {
    "T1": ("t1",),
    "T2*": ("t2",),
    "t2_echo": ("t2_echo",),
    "Qi": ("resonator_qi",),
    "Qc": ("resonator_qc",),
    "fr": ("resonator_freq_low", "resonator_freq_high"),
    "qubit_frequency": ("qubit_01_frequency",),
    "transition_frequency": ("qubit_01_frequency",),
    "pi_pulse_amplitude": ("pi_pulse_amp",),
}
QOI_ANALYSIS_ALIASES_LOWER = {
    key.lower(): value for key, value in QOI_ANALYSIS_ALIASES.items()
}


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 standard 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_mapping_keys(raw_record)
    raw_extra_tags = record.pop(EXTRA_TAGS_KEY, None)
    raw_extra_fields = record.pop(EXTRA_FIELDS_KEY, None)

    partner_extra_tags: dict[str, object] = {}
    for source_key, target_key in PARTNER_EXTRA_TAG_ALIASES.items():
        value = record.pop(source_key, None)
        if value is None:
            continue
        partner_extra_tags[target_key] = value

    merged_extra_tags: dict[str, str] = {}
    normalized_extra_tags = _as_string_map(raw_extra_tags)
    if normalized_extra_tags is not None:
        merged_extra_tags.update(normalized_extra_tags)

    normalized_partner_tags = _as_string_map(partner_extra_tags)
    if normalized_partner_tags is not None:
        merged_extra_tags.update(normalized_partner_tags)

    return (
        record,
        merged_extra_tags or None,
        _as_field_value_map(raw_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_mapping_keys(item))
        return records

    if isinstance(payload, Mapping):
        payload_map = _normalize_mapping_keys(payload)

        snapshot_records = _snapshot_payload_to_records(payload_map)
        if snapshot_records:
            return snapshot_records

        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
            record = _record_from_qoi_map_value(str(qoi_name), value)
            if record is not None:
                records.append(record)
        return records

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


def hybrid_qoi_records_from_payloads(
    *,
    snapshot_payload: Mapping[str, object],
    qoi_payloads: Sequence[Mapping[str, object]],
    dataset_tuid: str | None,
    preferred_elements: Sequence[str] | None = None,
) -> list[dict[str, object]]:
    """Build records by combining QOI-map values with snapshot metadata.

    The values come from ``quantities_of_interest`` payloads while record metadata
    (element, labels, units, run/device/cycle tags, condition) comes from snapshot.
    """
    snapshot_map = _normalize_mapping_keys(snapshot_payload)
    metadata_index = _snapshot_qoi_metadata_index(snapshot_map)
    if not metadata_index:
        return []

    qoi_values: dict[str, tuple[object, object | None]] = {}
    for raw_payload in qoi_payloads:
        payload_map = _normalize_mapping_keys(raw_payload)
        for qoi_name, value in payload_map.items():
            if qoi_name in IGNORED_QOI_MAP_KEYS:
                continue
            record = _record_from_qoi_map_value(qoi_name, value)
            if record is None:
                continue
            nominal_value = record.get("nominal_value")
            if _is_missing_number(nominal_value):
                continue
            uncertainty = record.get("uncertainty")
            for standard_qoi in _resolve_to_standard_qois(qoi_name):
                if standard_qoi not in SCQT_TRANSMON_QOI_KEYS:
                    continue
                qoi_values[standard_qoi] = (nominal_value, uncertainty)

    if not qoi_values:
        return []

    preferred_element_set = {
        item.strip()
        for item in (preferred_elements or [])
        if isinstance(item, str) and item.strip()
    }

    merged_records: list[dict[str, object]] = []
    seen: set[tuple[str, str, str]] = set()

    for standard_qoi, (nominal_value, uncertainty) in qoi_values.items():
        metadata_entries = metadata_index.get(standard_qoi, [])
        if not metadata_entries:
            continue

        matching_entries = metadata_entries
        if isinstance(dataset_tuid, str) and dataset_tuid.strip():
            matching_entries = [
                item for item in metadata_entries if item.get("tuid") == dataset_tuid
            ]
        if not matching_entries and preferred_element_set:
            matching_entries = [
                item
                for item in metadata_entries
                if item.get("element") in preferred_element_set
            ]
        if not matching_entries:
            matching_entries = metadata_entries

        for metadata in matching_entries:
            tuid = metadata.get("tuid")
            if isinstance(dataset_tuid, str) and dataset_tuid.strip():
                tuid = dataset_tuid.strip()
            if not isinstance(tuid, str) or not tuid.strip():
                continue

            element = metadata.get("element")
            if not isinstance(element, str) or not element.strip():
                continue

            dedupe_key = (standard_qoi, element.strip(), tuid.strip())
            if dedupe_key in seen:
                continue
            seen.add(dedupe_key)

            record: dict[str, object] = {
                "qoi": standard_qoi,
                "nominal_value": nominal_value,
                "tuid": tuid.strip(),
                "element": element.strip(),
            }

            for key in (
                "label",
                "unit",
                "element_label",
                "device_ID",
                "run_ID",
                "cycle_ID",
            ):
                value = metadata.get(key)
                if isinstance(value, str) and value.strip():
                    record[key] = value.strip()

            uncertainty_value = uncertainty
            if uncertainty_value is None:
                uncertainty_value = metadata.get("uncertainty", math.nan)
            if not _is_missing_number(uncertainty_value):
                record["uncertainty"] = uncertainty_value

            condition = metadata.get("condition")
            if isinstance(condition, str) and condition.strip():
                record["condition"] = condition.strip()

            merged_records.append(record)

    return merged_records


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

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

    Returns
    -------
    dict[str, object] | None
        Standard record dictionary containing at least ``qoi``, or ``None``
        when the payload entry is not a QOI value.
    """
    if isinstance(value, Mapping):
        value_map = _normalize_mapping_keys(value)
        if value_map.get("__dtype__") == "UFloat":
            return {
                "qoi": qoi_name,
                "nominal_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]
        for source_key, target_key in PARTNER_EXTRA_TAG_ALIASES.items():
            if source_key in value_map:
                record[target_key] = value_map[source_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

    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return {"qoi": qoi_name, "nominal_value": value}
    return None


def _normalize_mapping_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 _resolve_to_standard_qois(qoi_name: str) -> tuple[str, ...]:
    """Resolve an analysis QOI key to standard SCQT transmon QOI names."""
    stripped = qoi_name.strip()
    if not stripped:
        return ()
    if stripped in SCQT_TRANSMON_QOI_KEYS:
        return (stripped,)

    direct = QOI_ANALYSIS_ALIASES.get(stripped)
    if direct is not None:
        return direct

    lowered = stripped.lower()
    if lowered in SCQT_TRANSMON_QOI_KEYS:
        return (lowered,)

    alias = QOI_ANALYSIS_ALIASES_LOWER.get(lowered)
    if alias is not None:
        return alias
    return ()


def _snapshot_payload_to_records(
    payload_map: Mapping[str, object],
) -> list[dict[str, object]]:
    """Extract QOI records from a Quantify snapshot payload.

    Parameters
    ----------
    payload_map : Mapping[str, object]
        Parsed JSON mapping.

    Returns
    -------
    list[dict[str, object]]
        Standard records. Empty list when payload is not a snapshot.
    """
    raw_instruments = payload_map.get(SNAPSHOT_INSTRUMENTS_KEY)
    if not isinstance(raw_instruments, Mapping):
        return []

    instruments = _normalize_mapping_keys(raw_instruments)
    shared_tags = _snapshot_shared_tags(instruments)

    records: list[dict[str, object]] = []
    for element_name, raw_instrument in instruments.items():
        if not isinstance(raw_instrument, Mapping):
            continue

        instrument = _normalize_mapping_keys(raw_instrument)
        raw_submodules = instrument.get(SNAPSHOT_SUBMODULES_KEY)
        if not isinstance(raw_submodules, Mapping):
            continue

        submodules = _normalize_mapping_keys(raw_submodules)
        raw_qoi_module = submodules.get(SNAPSHOT_QOI_MODULE_KEY)
        if not isinstance(raw_qoi_module, Mapping):
            continue

        qoi_module = _normalize_mapping_keys(raw_qoi_module)
        raw_qoi_params = qoi_module.get(SNAPSHOT_PARAMETERS_KEY)
        if not isinstance(raw_qoi_params, Mapping):
            continue

        element_label = _as_optional_string(instrument.get("label"))
        for qoi_name, raw_qoi_param in _normalize_mapping_keys(raw_qoi_params).items():
            standard_names = _resolve_to_standard_qois(qoi_name)
            if len(standard_names) != 1:
                continue
            standard_qoi = standard_names[0]
            if standard_qoi not in SCQT_TRANSMON_QOI_KEYS:
                continue

            if not isinstance(raw_qoi_param, Mapping):
                continue

            qoi_param = _normalize_mapping_keys(raw_qoi_param)
            nominal_value = qoi_param.get("value")
            if _is_missing_number(nominal_value):
                continue

            raw_metadata = qoi_param.get("metadata")
            metadata = (
                _normalize_mapping_keys(raw_metadata)
                if isinstance(raw_metadata, Mapping)
                else {}
            )

            tuid = metadata.get("TUID")
            if not isinstance(tuid, str) or not tuid.strip():
                # Match the source behavior: points without TUID are not written.
                continue

            record: dict[str, object] = {
                "qoi": standard_qoi,
                "nominal_value": nominal_value,
                "tuid": tuid.strip(),
                "element": element_name,
                **shared_tags,
            }

            label = _as_optional_string(qoi_param.get("label"))
            if label is not None:
                record["label"] = label

            unit = _as_optional_string(qoi_param.get("unit"))
            if unit is not None:
                record["unit"] = unit

            if element_label is not None:
                record["element_label"] = element_label

            uncertainty = metadata.get("std_dev")
            if uncertainty is None:
                uncertainty = math.nan
            if not _is_missing_number(uncertainty):
                record["uncertainty"] = uncertainty

            condition = _as_optional_string(metadata.get("condition"))
            if condition is not None:
                record["condition"] = condition

            records.append(record)

    return records


def _snapshot_qoi_metadata_index(
    payload_map: Mapping[str, object],
) -> dict[str, list[dict[str, object]]]:
    """Index standard snapshot QoI metadata by QoI name."""
    raw_instruments = payload_map.get(SNAPSHOT_INSTRUMENTS_KEY)
    if not isinstance(raw_instruments, Mapping):
        return {}

    instruments = _normalize_mapping_keys(raw_instruments)
    shared_tags = _snapshot_shared_tags(instruments)

    index: dict[str, list[dict[str, object]]] = {}
    for element_name, raw_instrument in instruments.items():
        if not isinstance(raw_instrument, Mapping):
            continue

        instrument = _normalize_mapping_keys(raw_instrument)
        raw_submodules = instrument.get(SNAPSHOT_SUBMODULES_KEY)
        if not isinstance(raw_submodules, Mapping):
            continue

        submodules = _normalize_mapping_keys(raw_submodules)
        raw_qoi_module = submodules.get(SNAPSHOT_QOI_MODULE_KEY)
        if not isinstance(raw_qoi_module, Mapping):
            continue

        qoi_module = _normalize_mapping_keys(raw_qoi_module)
        raw_qoi_params = qoi_module.get(SNAPSHOT_PARAMETERS_KEY)
        if not isinstance(raw_qoi_params, Mapping):
            continue

        element_label = _as_optional_string(instrument.get("label"))
        for qoi_name, raw_qoi_param in _normalize_mapping_keys(raw_qoi_params).items():
            standard_names = _resolve_to_standard_qois(qoi_name)
            if len(standard_names) != 1:
                continue
            standard_qoi = standard_names[0]
            if standard_qoi not in SCQT_TRANSMON_QOI_KEYS:
                continue
            if not isinstance(raw_qoi_param, Mapping):
                continue

            qoi_param = _normalize_mapping_keys(raw_qoi_param)
            raw_metadata = qoi_param.get("metadata")
            metadata_map = (
                _normalize_mapping_keys(raw_metadata)
                if isinstance(raw_metadata, Mapping)
                else {}
            )

            metadata: dict[str, object] = {
                "element": element_name,
                **shared_tags,
            }

            label = _as_optional_string(qoi_param.get("label"))
            if label is not None:
                metadata["label"] = label

            unit = _as_optional_string(qoi_param.get("unit"))
            if unit is not None:
                metadata["unit"] = unit

            if element_label is not None:
                metadata["element_label"] = element_label

            tuid = _as_optional_string(metadata_map.get("TUID"))
            if tuid is not None:
                metadata["tuid"] = tuid

            uncertainty = metadata_map.get("std_dev")
            if uncertainty is None:
                uncertainty = math.nan
            if not _is_missing_number(uncertainty):
                metadata["uncertainty"] = uncertainty

            condition = _as_optional_string(metadata_map.get("condition"))
            if condition is not None:
                metadata["condition"] = condition

            index.setdefault(standard_qoi, []).append(metadata)

    return index


def _snapshot_shared_tags(instruments: Mapping[str, object]) -> dict[str, str]:
    """Extract device-level tags from snapshot instruments."""
    raw_device = instruments.get(SNAPSHOT_QUANTUM_DEVICE_KEY)
    if not isinstance(raw_device, Mapping):
        return {}

    quantum_device = _normalize_mapping_keys(raw_device)
    raw_params = quantum_device.get(SNAPSHOT_PARAMETERS_KEY)
    if not isinstance(raw_params, Mapping):
        return {}

    params = _normalize_mapping_keys(raw_params)
    shared_tags: dict[str, str] = {}
    for key in ("device_ID", "run_ID", "cycle_ID"):
        value = _snapshot_parameter_value(params.get(key))
        clean_value = _as_optional_string(value)
        if clean_value is not None:
            shared_tags[key] = clean_value
    return shared_tags


def _snapshot_parameter_value(raw_param: object) -> object | None:
    """Return the ``value`` field from a snapshot parameter mapping."""
    if not isinstance(raw_param, Mapping):
        return None
    param = _normalize_mapping_keys(raw_param)
    return param.get("value")


def _as_optional_string(value: object) -> str | None:
    """Normalize a non-empty string or return ``None``."""
    if not isinstance(value, str):
        return None
    stripped = value.strip()
    if not stripped:
        return None
    return stripped


def _is_missing_number(value: object) -> bool:
    """Return ``True`` for missing/NaN numeric values."""
    return value is None or (isinstance(value, float) and math.isnan(value))


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
