"""
Band / RAT-column mapper for raster studies.

Raster-aware analyses receive their inputs as a pandas DataFrame
(the raster's RAT), so the mapping from source columns to engine
variables is structurally identical to the vector case handled by
ParameterMapper. BandMapper is a thin adapter that lets the raster
Create-Analysis dialog speak in band-aware vocabulary ("band 1",
"impervious_pct from the land-cover RAT") and translates that into
the existing VariableMapping shape so all the type-conversion and
validation code from ParameterMapper is reused.

The mapping inputs come from the Study.input_layer_refs that the
raster Create-Study-Area dialog populated; each LayerRef carries
`bands_to_carry` (which RasterOverlay stores in the RAT under the
input's name) and `rat_columns_to_carry` (stored as
"<input_name>__<column>" — see RasterOverlay._enrich_with_input_rats).

This module is pure-Python and unit-testable without QGIS.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, List, Optional

import pandas as pd

from .parameter_mapper import (
    ParameterMapper,
    ParameterMappingConfig,
    VariableMapping,
    VariableType,
)


@dataclass
class BandMapping:
    """Maps one RAT column (a carried band value or RAT attribute) to
    a model variable.

    `source_input`: the input raster's display name (matches the
        Study.input_layer_refs[*].name the user picked in the dialog)
    `source_kind`:  "band" — the input's own pixel value column
                    (stored in the RAT as just `<source_input>`)
                    OR "rat_column" — a column joined in from the
                    input's RAT (stored as `<source_input>__<column>`)
    `source_column_or_band`: For "band" mapping, the band index
        (informational only — the RAT column is the input name).
        For "rat_column" mapping, the source RAT column name.
    `target_variable`: the engine model variable to bind to.
    """

    source_input: str
    source_kind: str  # "band" | "rat_column"
    source_column_or_band: str
    target_variable: str
    variable_type: VariableType = VariableType.DOUBLE
    required: bool = True

    BAND = "band"
    RAT_COLUMN = "rat_column"

    def resolve_source_column(self) -> str:
        """Compute the RAT column name to look up in the input DataFrame.

        Mirrors the naming convention used by
        RasterOverlay._enrich_with_input_rats so the RAT shape and the
        mapper agree on column names.
        """
        if self.source_kind == self.BAND:
            return self.source_input
        if self.source_kind == self.RAT_COLUMN:
            return f"{self.source_input}__{self.source_column_or_band}"
        raise ValueError(
            f"Unknown BandMapping.source_kind '{self.source_kind}'. "
            f"Use BandMapping.BAND or BandMapping.RAT_COLUMN."
        )

    def to_variable_mapping(self) -> VariableMapping:
        """Lower into the ParameterMapper-native VariableMapping shape."""
        return VariableMapping(
            source_column=self.resolve_source_column(),
            target_variable=self.target_variable,
            variable_type=self.variable_type,
            required=self.required,
            type_conversion=True,
        )


class BandMapper:
    """Maps a raster study's enriched RAT to engine model variables.

    Construct with a list of BandMapping, then call `map_dataframe(rat_df)`
    to produce the mapped DataFrame the engine bridges consume.
    Internally delegates to ParameterMapper so the type-conversion and
    missing-column-validation behavior matches the vector path exactly.
    """

    def __init__(
        self,
        *,
        workflow_step_id: str,
        model_name: str,
        mappings: Iterable[BandMapping],
    ):
        self._band_mappings: List[BandMapping] = list(mappings)
        self._underlying_config = ParameterMappingConfig(
            workflow_step_id=workflow_step_id,
            model_name=model_name,
            mappings=[bm.to_variable_mapping() for bm in self._band_mappings],
        )
        self._mapper = ParameterMapper(self._underlying_config)

    def map_dataframe(self, rat_df: pd.DataFrame) -> pd.DataFrame:
        """Apply the band mappings to a raster's enriched RAT."""
        # ParameterMapper API may differ between builds; the project
        # uses .map_dataframe in some places. Probe for it lazily.
        if hasattr(self._mapper, "map_dataframe"):
            return self._mapper.map_dataframe(rat_df)
        if hasattr(self._mapper, "apply"):
            return self._mapper.apply(rat_df)
        # Fallback: build the mapped df by hand using the same lookup
        # rules ParameterMapper uses (case-insensitive column match).
        return self._fallback_map_dataframe(rat_df)

    def _fallback_map_dataframe(self, rat_df: pd.DataFrame) -> pd.DataFrame:
        out = pd.DataFrame(index=rat_df.index)
        cols_lower = {col.lower(): col for col in rat_df.columns}
        missing_required: List[str] = []
        for bm in self._band_mappings:
            src = bm.resolve_source_column()
            actual = cols_lower.get(src.lower())
            if actual is None:
                if bm.required:
                    missing_required.append(src)
                continue
            out[bm.target_variable] = rat_df[actual]
        if missing_required:
            raise ValueError(
                f"BandMapper: required RAT columns missing from input: "
                f"{missing_required}. Available columns: {list(rat_df.columns)}."
            )
        return out

    @property
    def mappings(self) -> List[BandMapping]:
        return list(self._band_mappings)
