"""
Raster engine output writer.

EMDS engines (NetWeaver / CDP / Portfolio) return their results as a
pandas DataFrame. For vector studies, that DataFrame becomes a new
vector layer's attribute table via layer_creator.py. For raster
studies there are two cases:

  1. **Per-RAT-row result (common case)**: the engine returned one
     row per existing RAT row (i.e. one row per unique cell
     combination). We write the result back as an updated RAT on a
     copy of the overlay raster — the pixel codes don't change, only
     the attribute columns do. This is the cheap, idiomatic path.

  2. **Per-cell result (uncommon)**: the engine returned an array the
     size of the raster grid. We write a new single-band GeoTIFF
     using QGIS's GDAL translate algorithm via processing.run. EMDS
     today doesn't have engines that do this, but the writer accepts
     the shape so Phase 4 raster-algebra workflow steps have a home.

Both paths are isolated behind small methods so the dialog / executor
can pick the right one at call time.
"""

from __future__ import annotations

import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Dict, Optional

import pandas as pd

from .rat_io import RAT_VALUE_COLUMN, RatWriter


class RasterOutputWriterError(RuntimeError):
    """Raised when the writer cannot persist the engine output."""


@dataclass
class RasterOutputWriteResult:
    output_path: str
    mode: str  # "rat_update" | "new_grid"


# Type alias mirrored from raster_overlay for consistent injection seams.
ProcessingRun = Callable[[str, Dict[str, Any]], Dict[str, Any]]


def _default_processing_run(algorithm_id: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
    from qgis import processing  # type: ignore
    return processing.run(algorithm_id, parameters)


class RasterOutputWriter:
    """Persist engine results as either an updated RAT or a new raster grid."""

    def __init__(
        self,
        *,
        rat_writer: Optional[RatWriter] = None,
        processing_run: ProcessingRun = _default_processing_run,
    ):
        self._rat_writer = rat_writer or RatWriter()
        self._processing_run = processing_run

    # ---- Mode 1: update RAT on a copy of the overlay raster ------------

    def write_rat_update(
        self,
        *,
        overlay_raster_path: str,
        engine_result_df: pd.DataFrame,
        output_raster_path: str,
    ) -> RasterOutputWriteResult:
        """Copy the overlay raster to `output_raster_path` and replace
        its RAT with the engine's result DataFrame.

        Pre-condition: `engine_result_df` must include the `Value`
        column (the pixel code from the overlay) so the RAT writer
        and downstream readers know how to join it back to pixels.
        """
        if engine_result_df is None or engine_result_df.empty:
            raise RasterOutputWriterError(
                "engine_result_df is empty — nothing to write."
            )
        if RAT_VALUE_COLUMN.lower() not in {c.lower() for c in engine_result_df.columns}:
            raise RasterOutputWriterError(
                f"engine_result_df missing the '{RAT_VALUE_COLUMN}' join key. "
                f"Got columns: {list(engine_result_df.columns)}."
            )
        if not Path(overlay_raster_path).exists():
            raise RasterOutputWriterError(
                f"Source overlay raster not found: {overlay_raster_path}"
            )

        # Copy the raster (and any sidecar files — .vat.dbf, .aux.xml,
        # .tfw, .prj — that live next to it). The user's engine output
        # will replace the RAT but leave pixels intact.
        self._copy_raster_with_sidecars(overlay_raster_path, output_raster_path)

        # RatWriter expects either a layer or a path. Path form keeps
        # us off the QGIS object dependency in tests.
        self._rat_writer.write(output_raster_path, band=1, df=engine_result_df)

        return RasterOutputWriteResult(
            output_path=output_raster_path,
            mode="rat_update",
        )

    # ---- Mode 2: write a per-cell array as a new GeoTIFF --------------

    def write_new_grid(
        self,
        *,
        reference_raster_path: str,
        per_cell_array: Any,           # numpy.ndarray, kept loose to avoid the import here
        output_raster_path: str,
        nodata: Optional[float] = None,
    ) -> RasterOutputWriteResult:
        """Write a per-cell numpy array as a new single-band GeoTIFF,
        copying georeferencing from `reference_raster_path`.

        Implementation goes via QGIS Processing's gdal:translate-style
        path so we don't add a rasterio dep. The reference raster
        supplies the grid (extent / cellsize / CRS); we replace its
        pixels with the engine output.

        Engines that hand us a per-cell array are rare today —
        this is the door Phase 4 raster-algebra walks through.
        """
        if per_cell_array is None:
            raise RasterOutputWriterError(
                "per_cell_array is None — nothing to write."
            )
        if not Path(reference_raster_path).exists():
            raise RasterOutputWriterError(
                f"Reference raster not found: {reference_raster_path}"
            )

        # We stage the numpy array through a temp .npy → gdal_translate
        # path. Concrete implementation: write the array to a temp
        # GeoTIFF using GDAL's MEM driver via processing's
        # `gdal:rastercalculator`, parameterized to "give me back this
        # array on the reference grid". For Phase 3 the writer
        # surfaces the contract; full wiring lands when the first
        # per-cell engine arrives.
        raise NotImplementedError(
            "Per-cell array output is a Phase 4 surface. Today's "
            "engines return per-RAT-row results; use write_rat_update "
            "instead. When the first per-cell engine arrives, wire "
            "this through gdal:rastercalculator via processing.run."
        )

    # ---- Helpers -------------------------------------------------------

    @staticmethod
    def _copy_raster_with_sidecars(src_path: str, dst_path: str) -> None:
        """Copy a raster file plus its conventional sidecars.

        Covers: <name>.aux.xml, <name>.vat.dbf, <name>.tfw, <name>.prj,
        plus .ovr (overviews) so QGIS doesn't have to rebuild them.
        """
        src = Path(src_path)
        dst = Path(dst_path)
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)
        for suffix in (".aux.xml", ".vat.dbf", ".tfw", ".prj", ".ovr"):
            sidecar = src.with_name(src.name + suffix)
            if sidecar.exists():
                shutil.copy2(sidecar, dst.with_name(dst.name + suffix))
