"""
Raster Attribute Table (RAT) read / write helpers for EMDS raster studies.

A RAT is a small table attached to a categorical raster: each row
describes one pixel value (e.g. land-cover code 11 -> "Open Water",
impervious_pct=0). EMDS raster study areas live or die by RAT
fidelity, because the Create-Study-Area dialog asks the user which
RAT columns to carry forward into the overlay output, and the
engines (NetWeaver / CDP / Portfolio) consume those carried columns
as a pandas DataFrame.

Two storage paths are supported:

1. **Native** — `QgsRasterLayer.attributeTable(band)` (added in QGIS
   3.30). This is the modern path and is preferred when available.

2. **Sidecar `.vat.dbf`** — the legacy ESRI format that lives next
   to a `.tif` as `<name>.tif.vat.dbf`. We read these via QGIS's
   OGR vector layer driver, which opens DBFs natively, then pull
   columns into a DataFrame. This is what ArcGIS Pro produces, and
   the format the EMDS-for-ArcGIS-Pro raster studies emit. We also
   *write* this format for back-compat on older QGIS builds.

Both reader and writer return / accept a pandas DataFrame keyed on
the pixel value column (conventionally named `Value` — the ESRI
convention — but the reader detects it case-insensitively).

Module is lazy-imported by callers and lazily imports qgis.core
inside each method, so it can be unit-tested without QGIS via the
injection points in `read_dbf_via_callable` / `write_dbf_via_callable`.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable, List, Optional

import pandas as pd


# Conventional RAT key-column name. ESRI / ArcGIS uses "Value"; QGIS's
# native RAT API also uses this. Reader / writer do case-insensitive
# lookup so quirky inputs survive a round-trip.
RAT_VALUE_COLUMN = "Value"


class RatNotAvailableError(RuntimeError):
    """Raised when a raster has no RAT and one was required."""


def _sidecar_dbf_path(raster_source: str) -> Path:
    """Compute the ESRI sidecar RAT path for a raster source string.

    ArcGIS Pro / GDAL write `<raster>.vat.dbf` next to a GeoTIFF. We
    accept either `foo.tif` (→ `foo.tif.vat.dbf`) or already-suffixed
    `foo.tif.vat.dbf` paths.
    """
    p = Path(raster_source)
    if p.name.lower().endswith(".vat.dbf"):
        return p
    return p.with_name(p.name + ".vat.dbf")


class RatReader:
    """Read a raster's attribute table into a pandas DataFrame.

    Usage:
        df = RatReader().read(qgs_raster_layer, band=1)

    Resolution order:
      1. Native `QgsRasterLayer.attributeTable(band)` if QGIS exposes
         it and the layer has one populated.
      2. Sidecar `<source>.vat.dbf` opened as a QGIS OGR vector
         layer; column values pulled into a DataFrame.

    Returns None if neither path yields a RAT — callers decide
    whether that's an error (RatNotAvailableError) or just means
    "this raster is single-value, treat the band value itself as
    the only carried column".
    """

    def read(self, raster_layer: Any, band: int = 1) -> Optional[pd.DataFrame]:
        df = self._read_native(raster_layer, band)
        if df is not None:
            return df
        source = getattr(raster_layer, "source", lambda: "")()
        if not source:
            return None
        return self._read_sidecar_dbf(source)

    # ---- Native QgsRasterAttributeTable path ---------------------------

    @staticmethod
    def _read_native(raster_layer: Any, band: int) -> Optional[pd.DataFrame]:
        # `attributeTable` is the public accessor on QgsRasterLayer
        # in QGIS ≥3.30. Older builds raise AttributeError; some
        # newer builds return None when the layer has no RAT loaded.
        try:
            rat = raster_layer.attributeTable(band)
        except AttributeError:
            return None
        except Exception:
            return None
        if rat is None:
            return None
        try:
            fields = list(rat.fields())  # list of QgsField
            column_names = [field.name() for field in fields]
            rows: List[dict] = []
            for raw_row in rat.data():
                row = {}
                for index, column_name in enumerate(column_names):
                    try:
                        row[column_name] = raw_row[index]
                    except Exception:
                        row[column_name] = None
                rows.append(row)
            return pd.DataFrame(rows, columns=column_names) if column_names else None
        except Exception:
            # Native API exists but doesn't behave as expected on this
            # build (likely an empty / unsupported provider). Fall
            # through so the caller can try the sidecar path.
            return None

    # ---- ESRI sidecar `.vat.dbf` path ----------------------------------

    def _read_sidecar_dbf(self, raster_source: str) -> Optional[pd.DataFrame]:
        dbf_path = _sidecar_dbf_path(raster_source)
        if not dbf_path.exists():
            return None
        return self.read_dbf_via_callable(
            str(dbf_path),
            opener=self._default_dbf_opener,
        )

    @staticmethod
    def read_dbf_via_callable(
        dbf_path: str,
        opener: Callable[[str], Optional[pd.DataFrame]],
    ) -> Optional[pd.DataFrame]:
        """Test-injection seam: opener takes a path, returns a DataFrame.

        Production wires `_default_dbf_opener` (QGIS OGR-vector). Tests
        wire a fake opener so the reader can be exercised without
        QGIS. Kept as a static method so the test never touches QGIS
        imports.
        """
        return opener(dbf_path)

    @staticmethod
    def _default_dbf_opener(dbf_path: str) -> Optional[pd.DataFrame]:
        """Open a DBF via QGIS's OGR vector driver and pull rows into pandas."""
        try:
            from qgis.core import QgsVectorLayer  # type: ignore
        except ImportError:
            return None
        layer = QgsVectorLayer(dbf_path, Path(dbf_path).stem, "ogr")
        if not layer.isValid():
            return None
        column_names = [field.name() for field in layer.fields()]
        rows = []
        for feature in layer.getFeatures():
            row = {name: feature[name] for name in column_names}
            rows.append(row)
        if not rows:
            return pd.DataFrame(columns=column_names) if column_names else None
        return pd.DataFrame(rows, columns=column_names)


class RatWriter:
    """Write a pandas DataFrame back as a raster attribute table.

    Preferred path: native `QgsRasterLayer.setAttributeTable(band, rat)`
    on QGIS ≥3.30 + the RasterAttributeTable provider. Fallback: write
    the sidecar `.vat.dbf` next to the raster file.

    DataFrame must contain a value column (case-insensitively `Value`)
    — that's the pixel-value join key for the engines downstream.
    """

    def write(
        self,
        raster_layer_or_path: Any,
        band: int,
        df: pd.DataFrame,
    ) -> None:
        if df is None or df.empty:
            return
        self._validate_value_column(df)

        if self._write_native(raster_layer_or_path, band, df):
            return

        # Fallback: write sidecar .vat.dbf alongside the raster file.
        source = self._resolve_raster_source(raster_layer_or_path)
        if not source:
            raise RatNotAvailableError(
                "Cannot write RAT: native QgsRasterAttributeTable is "
                "unavailable on this QGIS build and the raster has no "
                "file-system source for a sidecar .vat.dbf."
            )
        self._write_sidecar_dbf(source, df)

    @staticmethod
    def _validate_value_column(df: pd.DataFrame) -> None:
        lowercase = {col.lower(): col for col in df.columns}
        if RAT_VALUE_COLUMN.lower() not in lowercase:
            raise ValueError(
                f"RAT DataFrame must contain a '{RAT_VALUE_COLUMN}' column "
                f"(case-insensitive). Got: {list(df.columns)}"
            )

    @staticmethod
    def _resolve_raster_source(raster_layer_or_path: Any) -> str:
        if isinstance(raster_layer_or_path, (str, Path)):
            return str(raster_layer_or_path)
        source_attr = getattr(raster_layer_or_path, "source", None)
        if callable(source_attr):
            try:
                return source_attr()
            except Exception:
                return ""
        return ""

    @staticmethod
    def _write_native(
        raster_layer_or_path: Any,
        band: int,
        df: pd.DataFrame,
    ) -> bool:
        """Try the native path; return True on success, False to fall back."""
        # Only available when we have a QgsRasterLayer instance (not a
        # string path), and only on QGIS builds that ship the native
        # RAT API. We try-import inside so this module imports cleanly
        # without QGIS.
        if isinstance(raster_layer_or_path, (str, Path)):
            return False
        try:
            from qgis.core import (  # type: ignore
                QgsField,
                QgsRasterAttributeTable,
            )
        except ImportError:
            return False
        except Exception:
            return False
        try:
            rat = QgsRasterAttributeTable()
            for column_name in df.columns:
                # Field type detection is intentionally loose — Int,
                # Double, String cover the EMDS use cases.
                series = df[column_name]
                if pd.api.types.is_integer_dtype(series):
                    field_type = QgsField.Int
                elif pd.api.types.is_float_dtype(series):
                    field_type = QgsField.Double
                else:
                    field_type = QgsField.String
                rat.appendField(column_name, field_type)
            for record in df.to_dict(orient="records"):
                rat.appendRow([record[col] for col in df.columns])
            raster_layer_or_path.setAttributeTable(band, rat)
            return True
        except AttributeError:
            return False
        except Exception:
            return False

    @staticmethod
    def _write_sidecar_dbf(raster_source: str, df: pd.DataFrame) -> None:
        """Write DataFrame to <raster>.vat.dbf via QGIS's OGR writer."""
        target = _sidecar_dbf_path(raster_source)
        try:
            from qgis.core import (  # type: ignore
                QgsVectorFileWriter,
                QgsVectorLayer,
                QgsField,
                QgsFeature,
            )
            from qgis.PyQt.QtCore import QVariant  # type: ignore
        except ImportError as exc:
            raise RatNotAvailableError(
                f"Cannot write sidecar .vat.dbf: QGIS not available ({exc})."
            )

        # Build an in-memory non-spatial vector layer, then write it
        # out via QgsVectorFileWriter with the ESRI DBF driver.
        # `None` geometry type makes this a non-spatial table.
        layer = QgsVectorLayer("None", "rat", "memory")
        provider = layer.dataProvider()
        qgis_fields = []
        for column_name in df.columns:
            series = df[column_name]
            if pd.api.types.is_integer_dtype(series):
                qt_type = QVariant.LongLong
            elif pd.api.types.is_float_dtype(series):
                qt_type = QVariant.Double
            else:
                qt_type = QVariant.String
            qgis_fields.append(QgsField(column_name, qt_type))
        provider.addAttributes(qgis_fields)
        layer.updateFields()

        features = []
        for record in df.to_dict(orient="records"):
            feat = QgsFeature(layer.fields())
            for col in df.columns:
                feat.setAttribute(col, record[col])
            features.append(feat)
        provider.addFeatures(features)

        # writeAsVectorFormat is deprecated but still works across the
        # 3.x line; the V2/V3 replacements moved across builds. Stick
        # with V2 for now (the codebase uses it elsewhere).
        save_options = None
        try:
            save_options = QgsVectorFileWriter.SaveVectorOptions()
            save_options.driverName = "ESRI Shapefile"  # writes .dbf
            save_options.fileEncoding = "UTF-8"
            QgsVectorFileWriter.writeAsVectorFormatV2(
                layer,
                str(target.with_suffix("")),  # writer adds .dbf
                layer.transformContext()
                    if hasattr(layer, "transformContext")
                    else None,
                save_options,
            )
        except Exception:
            # Last-resort: writeAsVectorFormat (no V2 / V3 suffix) —
            # exists on older builds.
            QgsVectorFileWriter.writeAsVectorFormat(  # type: ignore[attr-defined]
                layer,
                str(target.with_suffix("")),
                "UTF-8",
                layer.crs() if hasattr(layer, "crs") else None,
                "ESRI Shapefile",
            )
