from __future__ import annotations

from typing import Any, Dict

from .cli_api import (
    BondAngleDistributionPlan,
    BondLengthDistributionPlan,
    ChainRgPlan,
    ConductivityPlan,
    ContourLengthPlan,
    DielectricPlan,
    DipoleAlignmentPlan,
    EndToEndPlan,
    EquipartitionPlan,
    HbondPlan,
    IonPairCorrelationPlan,
    MsdPlan,
    PersistenceLengthPlan,
    RdfPlan,
    RgPlan,
    RmsdPlan,
    RotAcfPlan,
    StructureFactorPlan,
    System,
    WaterCountPlan,
    _select,
)
from .cli_utils import _as_tuple, _pick, _resolve_charges, _resolve_group_types


def _build_rg(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "rg.selection")
    kwargs = _pick(spec, ["mass_weighted"])
    return RgPlan(sel, **kwargs)


def _build_rmsd(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "rmsd.selection")
    kwargs = _pick(spec, ["reference", "align"])
    return RmsdPlan(sel, **kwargs)


def _build_msd(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "msd.selection")
    group_by = spec.get("group_by", "resid")
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(
        spec,
        [
            "axis",
            "length_scale",
            "frame_decimation",
            "dt_decimation",
            "time_binning",
            "lag_mode",
            "max_lag",
            "memory_budget_bytes",
            "multi_tau_m",
            "multi_tau_levels",
        ],
    )
    if "axis" in kwargs:
        kwargs["axis"] = _as_tuple(kwargs["axis"], 3, "axis")
    if "frame_decimation" in kwargs:
        kwargs["frame_decimation"] = _as_tuple(kwargs["frame_decimation"], 2, "frame_decimation")
    if "dt_decimation" in kwargs:
        kwargs["dt_decimation"] = _as_tuple(kwargs["dt_decimation"], 4, "dt_decimation")
    if "time_binning" in kwargs:
        kwargs["time_binning"] = _as_tuple(kwargs["time_binning"], 2, "time_binning")
    if group_types is not None:
        kwargs["group_types"] = group_types
    return MsdPlan(sel, group_by=group_by, **kwargs)


def _build_rotacf(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "rotacf.selection")
    group_by = spec.get("group_by", "resid")
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(
        spec,
        [
            "orientation",
            "p2_legendre",
            "length_scale",
            "frame_decimation",
            "dt_decimation",
            "time_binning",
            "lag_mode",
            "max_lag",
            "memory_budget_bytes",
            "multi_tau_m",
            "multi_tau_levels",
        ],
    )
    if "orientation" in kwargs:
        orient = kwargs["orientation"]
        if not isinstance(orient, (list, tuple)) or len(orient) not in (2, 3):
            raise ValueError("rotacf.orientation must be length 2 or 3")
    if "frame_decimation" in kwargs:
        kwargs["frame_decimation"] = _as_tuple(kwargs["frame_decimation"], 2, "frame_decimation")
    if "dt_decimation" in kwargs:
        kwargs["dt_decimation"] = _as_tuple(kwargs["dt_decimation"], 4, "dt_decimation")
    if "time_binning" in kwargs:
        kwargs["time_binning"] = _as_tuple(kwargs["time_binning"], 2, "time_binning")
    if group_types is not None:
        kwargs["group_types"] = group_types
    return RotAcfPlan(sel, group_by=group_by, **kwargs)


def _build_conductivity(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "conductivity.selection")
    group_by = spec.get("group_by", "resid")
    charges_spec = spec.get("charges")
    if charges_spec is None:
        raise ValueError("conductivity.charges is required")
    charges = _resolve_charges(system, charges_spec)
    temperature = spec.get("temperature")
    if temperature is None:
        raise ValueError("conductivity.temperature is required")
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(
        spec,
        [
            "transference",
            "length_scale",
            "frame_decimation",
            "dt_decimation",
            "time_binning",
            "lag_mode",
            "max_lag",
            "memory_budget_bytes",
            "multi_tau_m",
            "multi_tau_levels",
        ],
    )
    if "frame_decimation" in kwargs:
        kwargs["frame_decimation"] = _as_tuple(kwargs["frame_decimation"], 2, "frame_decimation")
    if "dt_decimation" in kwargs:
        kwargs["dt_decimation"] = _as_tuple(kwargs["dt_decimation"], 4, "dt_decimation")
    if "time_binning" in kwargs:
        kwargs["time_binning"] = _as_tuple(kwargs["time_binning"], 2, "time_binning")
    if group_types is not None:
        kwargs["group_types"] = group_types
    return ConductivityPlan(sel, charges, temperature, group_by=group_by, **kwargs)


def _build_dielectric(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "dielectric.selection")
    group_by = spec.get("group_by", "resid")
    charges_spec = spec.get("charges")
    if charges_spec is None:
        raise ValueError("dielectric.charges is required")
    charges = _resolve_charges(system, charges_spec)
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(spec, ["length_scale"])
    if group_types is not None:
        kwargs["group_types"] = group_types
    return DielectricPlan(sel, charges, group_by=group_by, **kwargs)


def _build_dipole_alignment(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "dipole_alignment.selection")
    group_by = spec.get("group_by", "resid")
    charges_spec = spec.get("charges")
    if charges_spec is None:
        raise ValueError("dipole_alignment.charges is required")
    charges = _resolve_charges(system, charges_spec)
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(spec, ["length_scale"])
    if group_types is not None:
        kwargs["group_types"] = group_types
    return DipoleAlignmentPlan(sel, charges, group_by=group_by, **kwargs)


def _build_ion_pair(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "ion_pair_correlation.selection")
    group_by = spec.get("group_by", "resid")
    rclust_cat = spec.get("rclust_cat")
    rclust_ani = spec.get("rclust_ani")
    if rclust_cat is None or rclust_ani is None:
        raise ValueError("ion_pair_correlation.rclust_cat and rclust_ani are required")
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(
        spec,
        [
            "cation_type",
            "anion_type",
            "max_cluster",
            "length_scale",
            "lag_mode",
            "max_lag",
            "memory_budget_bytes",
            "multi_tau_m",
            "multi_tau_levels",
        ],
    )
    if group_types is not None:
        kwargs["group_types"] = group_types
    return IonPairCorrelationPlan(sel, rclust_cat, rclust_ani, group_by=group_by, **kwargs)


def _build_structure_factor(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "structure_factor.selection")
    bins = spec.get("bins")
    r_max = spec.get("r_max")
    q_bins = spec.get("q_bins")
    q_max = spec.get("q_max")
    if None in (bins, r_max, q_bins, q_max):
        raise ValueError("structure_factor requires bins, r_max, q_bins, q_max")
    kwargs = _pick(spec, ["pbc", "length_scale"])
    return StructureFactorPlan(sel, bins, r_max, q_bins, q_max, **kwargs)


def _build_water_count(system: System, spec: Dict[str, Any]):
    water_sel = _select(system, spec.get("water_selection"), "water_count.water_selection")
    center_sel = _select(system, spec.get("center_selection"), "water_count.center_selection")
    box_unit = spec.get("box_unit")
    region_size = spec.get("region_size")
    if box_unit is None or region_size is None:
        raise ValueError("water_count requires box_unit and region_size")
    kwargs = _pick(spec, ["shift", "length_scale"])
    kwargs["box_unit"] = _as_tuple(box_unit, 3, "box_unit")
    kwargs["region_size"] = _as_tuple(region_size, 3, "region_size")
    if "shift" in kwargs:
        kwargs["shift"] = _as_tuple(kwargs["shift"], 3, "shift")
    return WaterCountPlan(water_sel, center_sel, **kwargs)


def _build_equipartition(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "equipartition.selection")
    group_by = spec.get("group_by", "resid")
    group_types = _resolve_group_types(system, sel, group_by, spec.get("group_types"))
    kwargs = _pick(spec, ["velocity_scale", "length_scale"])
    if group_types is not None:
        kwargs["group_types"] = group_types
    return EquipartitionPlan(sel, group_by=group_by, **kwargs)


def _build_hbond(system: System, spec: Dict[str, Any]):
    donors = _select(system, spec.get("donors"), "hbond.donors")
    acceptors = _select(system, spec.get("acceptors"), "hbond.acceptors")
    dist_cutoff = spec.get("dist_cutoff")
    if dist_cutoff is None:
        raise ValueError("hbond.dist_cutoff is required")
    hydrogens_expr = spec.get("hydrogens")
    angle_cutoff = spec.get("angle_cutoff")
    if hydrogens_expr:
        hydrogens = _select(system, hydrogens_expr, "hbond.hydrogens")
        return HbondPlan(donors, acceptors, dist_cutoff, hydrogens=hydrogens, angle_cutoff=angle_cutoff)
    return HbondPlan(donors, acceptors, dist_cutoff)


def _build_rdf(system: System, spec: Dict[str, Any]):
    sel_a = _select(system, spec.get("sel_a"), "rdf.sel_a")
    sel_b = _select(system, spec.get("sel_b"), "rdf.sel_b")
    bins = spec.get("bins")
    r_max = spec.get("r_max")
    if bins is None or r_max is None:
        raise ValueError("rdf requires bins and r_max")
    kwargs = _pick(spec, ["pbc"])
    return RdfPlan(sel_a, sel_b, bins, r_max, **kwargs)


def _build_end_to_end(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "end_to_end.selection")
    return EndToEndPlan(sel)


def _build_contour_length(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "contour_length.selection")
    return ContourLengthPlan(sel)


def _build_chain_rg(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "chain_rg.selection")
    return ChainRgPlan(sel)


def _build_bond_length(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "bond_length_distribution.selection")
    bins = spec.get("bins")
    r_max = spec.get("r_max")
    if bins is None or r_max is None:
        raise ValueError("bond_length_distribution requires bins and r_max")
    return BondLengthDistributionPlan(sel, bins, r_max)


def _build_bond_angle(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "bond_angle_distribution.selection")
    bins = spec.get("bins")
    if bins is None:
        raise ValueError("bond_angle_distribution requires bins")
    kwargs = _pick(spec, ["degrees"])
    return BondAngleDistributionPlan(sel, bins, **kwargs)


def _build_persistence(system: System, spec: Dict[str, Any]):
    sel = _select(system, spec.get("selection"), "persistence_length.selection")
    return PersistenceLengthPlan(sel)


PLAN_BUILDERS = {
    "rg": _build_rg,
    "rmsd": _build_rmsd,
    "msd": _build_msd,
    "rotacf": _build_rotacf,
    "conductivity": _build_conductivity,
    "dielectric": _build_dielectric,
    "dipole_alignment": _build_dipole_alignment,
    "ion_pair_correlation": _build_ion_pair,
    "structure_factor": _build_structure_factor,
    "water_count": _build_water_count,
    "equipartition": _build_equipartition,
    "hbond": _build_hbond,
    "rdf": _build_rdf,
    "end_to_end": _build_end_to_end,
    "contour_length": _build_contour_length,
    "chain_rg": _build_chain_rg,
    "bond_length_distribution": _build_bond_length,
    "bond_angle_distribution": _build_bond_angle,
    "persistence_length": _build_persistence,
}


CLI_TO_PLAN = {
    "rg": "rg",
    "rmsd": "rmsd",
    "msd": "msd",
    "rotacf": "rotacf",
    "conductivity": "conductivity",
    "dielectric": "dielectric",
    "dipole-alignment": "dipole_alignment",
    "ion-pair-correlation": "ion_pair_correlation",
    "structure-factor": "structure_factor",
    "water-count": "water_count",
    "equipartition": "equipartition",
    "hbond": "hbond",
    "rdf": "rdf",
    "end-to-end": "end_to_end",
    "contour-length": "contour_length",
    "chain-rg": "chain_rg",
    "bond-length-distribution": "bond_length_distribution",
    "bond-angle-distribution": "bond_angle_distribution",
    "persistence-length": "persistence_length",
}
