# -*- coding: utf-8 -*-
"""Tool 1: Load GEE composite + indices + temporal metrics as one stack."""

import os
import re
import tempfile
import traceback

from qgis.PyQt.QtCore import Qt, QDate, QSettings
from qgis.PyQt.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout, QGroupBox, QComboBox,
    QDateEdit, QSpinBox, QCheckBox, QPushButton, QLabel, QMessageBox,
    QApplication, QGridLayout, QScrollArea, QWidget, QInputDialog,
)
from qgis.gui import QgsMapLayerComboBox
from .polygon_tool import PolygonDrawTool
from qgis.core import (QgsMapLayerProxyModel, QgsGeometry, QgsProject, QgsRasterLayer,
    QgsTask, QgsApplication, QgsCoordinateTransform)

from .. import gee_utils as gu
from .help_texts import help_button

def _coerce(v):
    """Codes may be stored as int or string in the asset."""
    try:
        return int(v)
    except (TypeError, ValueError):
        return v


(AOI_CANVAS, AOI_ADMIN, AOI_ACTIVE, AOI_LAYER, AOI_ASSET,
 AOI_DRAW) = range(6)
SEN_S2, SEN_L8, SEN_L9, SEN_L89, SEN_S1, SEN_FUSION, SEN_EMBED = range(7)


class LoaderDialog(QDialog):
    GAMMA = 2.0        # fixed RGB brightness (applied, not shown)

    def __init__(self, iface, state, parent=None):
        super().__init__(parent)
        self.iface = iface
        self.state = state
        self.drawn_geom = None
        self.poly_tool = None
        self.setWindowTitle("1. Load GEE Data & Indices")
        self._build_ui()
        self.resize(1020, 700)

    # ------------------------------------------------------------ UI
    def _build_ui(self):
        root = QVBoxLayout(self)
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        content = QWidget()
        cols = QHBoxLayout(content)
        left = QVBoxLayout()
        right = QVBoxLayout()
        cols.addLayout(left, 1)
        cols.addLayout(right, 1)

        # ================= LEFT COLUMN =================
        gb = QGroupBox("Satellite data")
        f = QFormLayout(gb)
        self.cb_sensor = QComboBox()
        self.cb_sensor.addItems([
            "Sentinel-2 SR (10 m)",
            "Landsat 8 (30 m)",
            "Landsat 9 (30 m)",
            "Landsat 8 + 9 (30 m)",
            "Sentinel-1 GRD (10 m)",
            "Sentinel-2 + Sentinel-1 fusion (10 m)",
            "Google Satellite Embedding (annual, 64 bands, 10 m)",
        ])
        f.addRow("Sensor:", self.cb_sensor)

        self.de_start = QDateEdit(QDate.currentDate().addYears(-1))
        self.de_end = QDateEdit(QDate.currentDate())
        for de in (self.de_start, self.de_end):
            de.setCalendarPopup(True)
            de.setDisplayFormat("yyyy-MM-dd")
        self.sp_cloud = QSpinBox()
        self.sp_cloud.setRange(0, 100); self.sp_cloud.setValue(30)
        self.sp_cloud.setSuffix(" %")
        hdate = QHBoxLayout()
        hdate.addWidget(self.de_start, 1)
        hdate.addWidget(QLabel("End:"))
        hdate.addWidget(self.de_end, 1)
        hdate.addWidget(QLabel("Cloud:"))
        hdate.addWidget(self.sp_cloud)
        f.addRow("Start:", hdate)

        # year selector: only for Satellite Embedding
        self.sp_year = QSpinBox()
        self.sp_year.setRange(2017, QDate.currentDate().year())
        self.sp_year.setValue(QDate.currentDate().year() - 1)
        f.addRow("Embedding year:", self.sp_year)

        self.cb_comp = QComboBox()
        self.cb_comp.addItems(["Median", "Mean", "Min", "Max"])
        f.addRow("Composite:", self.cb_comp)

        left.addWidget(gb)

        gb2 = QGroupBox("2 - Area of interest")
        f2 = QFormLayout(gb2)
        self.cb_aoi = QComboBox()
        self.cb_aoi.addItems([
            "Current map canvas extent",
            "Administrative unit (bundled)",
            "Active layer - selected features",
            "Polygon layer (choose below)",
            "GEE asset (advanced)",
            "Draw polygon on canvas",
        ])
        self.cb_aoi.currentIndexChanged.connect(self._aoi_mode_changed)
        f2.addRow("AOI source:", self.cb_aoi)

        # cascading bundled admin units (offline, no GEE round trip)
        self.cb_prov = QComboBox()
        self.cb_dist = QComboBox()
        self.cb_comm = QComboBox()
        self.cb_prov.currentIndexChanged.connect(self._cascade_district)
        self.cb_dist.currentIndexChanged.connect(self._cascade_commune)
        f2.addRow("Province:", self.cb_prov)
        f2.addRow("District:", self.cb_dist)
        f2.addRow("Commune:", self.cb_comm)

        self.cb_layer = QgsMapLayerComboBox()
        self.cb_layer.setFilters(QgsMapLayerProxyModel.Filter.PolygonLayer)
        self.chk_selected = QCheckBox("Use only selected feature(s)")
        self.chk_selected.setChecked(True)
        f2.addRow("Polygon layer:", self.cb_layer)
        f2.addRow("", self.chk_selected)

        self.cb_asset = QComboBox(); self.cb_asset.setEditable(True)
        self.cb_asset.addItems(gu.AOI_ASSETS)
        self.btn_afields = QPushButton("Load fields")
        self.btn_afields.clicked.connect(self._load_asset_fields)
        h_asset = QHBoxLayout()
        h_asset.addWidget(self.cb_asset, 1); h_asset.addWidget(self.btn_afields)
        f2.addRow("Asset path:", h_asset)
        self.cb_afield = QComboBox()
        self.cb_afield.currentTextChanged.connect(self._load_asset_values)
        self.cb_avalue = QComboBox(); self.cb_avalue.setEditable(True)
        self.chk_whole_asset = QCheckBox("Use whole asset (no filter)")
        f2.addRow("Filter field:", self.cb_afield)
        f2.addRow("Filter value:", self.cb_avalue)
        f2.addRow("", self.chk_whole_asset)

        self.btn_draw = QPushButton("Draw polygon ...")
        self.btn_draw.clicked.connect(self._start_draw)
        self.lbl_rect = QLabel("No polygon drawn.")
        f2.addRow(self.btn_draw, self.lbl_rect)
        lbl_hint = QLabel("Left-click adds points; right-click or "
                          "double-click finishes.")
        lbl_hint.setStyleSheet("color:#888; font-size:10px;")
        lbl_hint.setWordWrap(True)
        f2.addRow("", lbl_hint)

        self.btn_preview_aoi = QPushButton("Preview AOI on map")
        self.btn_preview_aoi.clicked.connect(self.preview_aoi)
        f2.addRow(self.btn_preview_aoi)
        self.lbl_aoi = QLabel("")
        self.lbl_aoi.setWordWrap(True)
        self.lbl_aoi.setStyleSheet("color:#666;")
        f2.addRow(self.lbl_aoi)
        left.addWidget(gb2)

        gbb = QGroupBox("Original bands (tick only what you need)")
        vb = QVBoxLayout(gbb)
        hbb = QHBoxLayout()
        bb_all = QPushButton("All")
        bb_rgb = QPushButton("RGB only")
        bb_rgbn = QPushButton("RGB + NIR")
        bb_none = QPushButton("None")
        bb_all.clicked.connect(lambda: self._set_base_bands("all"))
        bb_rgb.clicked.connect(lambda: self._set_base_bands("rgb"))
        bb_rgbn.clicked.connect(lambda: self._set_base_bands("rgbn"))
        bb_none.clicked.connect(lambda: self._set_base_bands("none"))
        for w in (bb_all, bb_rgb, bb_rgbn, bb_none):
            hbb.addWidget(w)
        hbb.addStretch()
        vb.addLayout(hbb)
        self.base_grid = QGridLayout()
        vb.addLayout(self.base_grid)
        self.base_checks = {}
        left.addWidget(gbb)
        self.gb_bands = gbb
        left.addStretch(1)

        # ================= RIGHT COLUMN =================
        gb4 = QGroupBox("Temporal metrics && extra bands")
        g4 = QGridLayout(gb4)
        self.chk_pct = QCheckBox("P20 + P80 percentile composites (per band)")
        self.chk_std = QCheckBox("Per-band StdDev (temporal variability)")
        self.chk_ndvit = QCheckBox("Temporal NDVI (mean/std/p20/p80)")
        self.chk_lst = QCheckBox("Landsat LST (Kelvin)")
        self.chk_terrain = QCheckBox("Terrain: NASADEM elevation + slope")
        g4.addWidget(self.chk_pct, 0, 0)
        g4.addWidget(self.chk_std, 1, 0)
        g4.addWidget(self.chk_ndvit, 2, 0)
        g4.addWidget(self.chk_lst, 3, 0)
        g4.addWidget(self.chk_terrain, 4, 0)
        self.chk_important = QCheckBox(
            "Important bands preset (~30 bands, paper style) -\n"
            "ticks the needed metrics/indices and drops the rest")
        self.chk_important.toggled.connect(self._important_toggled)
        g4.addWidget(self.chk_important, 5, 0)
        right.addWidget(gb4)

        gb3 = QGroupBox("Spectral indices (added as extra bands)")
        v3 = QVBoxLayout(gb3)
        hb = QHBoxLayout()
        b_all = QPushButton("Select all"); b_none = QPushButton("Select none")
        b_all.clicked.connect(lambda: self._set_all_indices(True))
        b_none.clicked.connect(lambda: self._set_all_indices(False))
        hb.addWidget(b_all); hb.addWidget(b_none); hb.addStretch()
        v3.addLayout(hb)
        g3 = QGridLayout()
        self.idx_checks = {}
        all_names = gu.OPTICAL_INDICES + gu.S2_ONLY_INDICES + gu.SAR_INDICES
        default_idx = {"NDVI", "NDWI"}
        for i, name in enumerate(all_names):
            cb = QCheckBox(name)
            cb.setChecked(name in default_idx)
            self.idx_checks[name] = cb
            g3.addWidget(cb, i // 3, i % 3)
        v3.addLayout(g3)
        right.addWidget(gb3)

        self.chk_local = QCheckBox(
            "Load stack as local multiband GeoTIFF\n"
            "(all bands separate; big AOIs are auto-resampled)")
        self.chk_local.setChecked(True)
        right.addWidget(self.chk_local)
        self.chk_add_idx_layers = QCheckBox(
            "Also add each index as its own preview layer")
        right.addWidget(self.chk_add_idx_layers)
        right.addStretch(1)

        self.cb_sensor.currentIndexChanged.connect(self._sensor_changed)

        scroll.setWidget(content)
        root.addWidget(scroll)

        # ---- fixed bottom: buttons + status
        h = QHBoxLayout()
        self.btn_run = QPushButton("Build stack && load to canvas")
        self.btn_run.clicked.connect(self.run)
        self.btn_drive = QPushButton("Export to Drive")
        self.btn_drive.clicked.connect(self.export_drive)
        self.btn_asset = QPushButton("Export to Asset")
        self.btn_asset.clicked.connect(self.export_asset)
        self.btn_load_asset = QPushButton("Load stack from Asset ...")
        self.btn_load_asset.clicked.connect(self.load_asset)
        self.btn_reset = QPushButton("Reset")
        self.btn_reset.clicked.connect(self.reset_defaults)
        self.btn_close = QPushButton("Close")
        self.btn_close.clicked.connect(self.close)
        # Build stack is the PRIMARY action: emphasised and set apart
        self.btn_run.setMinimumHeight(34)
        self.btn_run.setStyleSheet("font-weight: bold;")
        h.addWidget(help_button("loader", self))
        h.addWidget(self.btn_run)
        h.addSpacing(38)                     # ~10 mm gap
        h.addWidget(self.btn_drive)
        h.addWidget(self.btn_asset)
        h.addWidget(self.btn_load_asset)
        h.addStretch()
        h.addWidget(self.btn_reset); h.addWidget(self.btn_close)
        root.addLayout(h)
        self.lbl_status = QLabel(""); self.lbl_status.setWordWrap(True)
        root.addWidget(self.lbl_status)

        self.cb_aoi.setCurrentIndex(AOI_ADMIN)
        self._sensor_changed()
        self._aoi_mode_changed(AOI_ADMIN)

    # ------------------------------------------------------------ presets
    def _important_toggled(self, on):
        if not on:
            return
        self.chk_pct.setChecked(True)
        self.chk_std.setChecked(True)
        self.chk_ndvit.setChecked(True)
        for n, cb in self.idx_checks.items():
            if cb.isEnabled():
                cb.setChecked(n in gu.IMPORTANT_INDEX_CHECKS)

    # ------------------------------------------------------------ sensor logic
    def _set_all_indices(self, checked):
        for cb in self.idx_checks.values():
            if cb.isEnabled():
                cb.setChecked(checked)

    def _base_band_list(self):
        i = self.cb_sensor.currentIndex()
        if i == SEN_S2:
            return gu.S2_BASE
        if i in (SEN_L8, SEN_L9, SEN_L89):
            return gu.LS_BASE
        if i == SEN_S1:
            return ["VV", "VH"]
        if i == SEN_FUSION:
            return gu.S2_BASE + ["VV", "VH"]
        return []                      # embedding: fixed 64 bands

    def _rebuild_base_checks(self):
        # clear old
        while self.base_grid.count():
            it = self.base_grid.takeAt(0)
            w = it.widget()
            if w:
                w.setParent(None)
        self.base_checks = {}
        bands = self._base_band_list()
        # default: RGB + NIR (plus both SAR polarisations when present)
        preset = {"red", "green", "blue", "nir", "VV", "VH"}
        for k, name in enumerate(bands):
            cb = QCheckBox(name)
            cb.setChecked(name in preset)
            self.base_checks[name] = cb
            self.base_grid.addWidget(cb, k // 3, k % 3)
        self.gb_bands.setVisible(bool(bands))
        self.gb_bands.setTitle(
            "Original bands (tick only what you need) - {} available"
            .format(len(bands)) if bands else "Original bands (n/a)")

    def _set_base_bands(self, mode):
        rgb = {"red", "green", "blue", "VV", "VH"}
        rgbn = rgb | {"nir"}
        for name, cb in self.base_checks.items():
            if mode == "all":
                cb.setChecked(True)
            elif mode == "none":
                cb.setChecked(False)
            elif mode == "rgbn":
                cb.setChecked(name in rgbn)
            else:                       # rgb
                cb.setChecked(name in rgb)

    def _sensor_changed(self, *_):
        i = self.cb_sensor.currentIndex()
        opt = i in (SEN_S2, SEN_L8, SEN_L9, SEN_L89, SEN_FUSION)
        re_ = i in (SEN_S2, SEN_FUSION)
        sar = i in (SEN_S1, SEN_FUSION)
        emb = i == SEN_EMBED
        default_idx = {"NDVI", "NDWI"}
        for n, cb in self.idx_checks.items():
            en = ((n in gu.OPTICAL_INDICES and opt)
                  or (n in gu.S2_ONLY_INDICES and re_)
                  or (n in gu.SAR_INDICES and sar))
            cb.setEnabled(en)
            cb.setChecked(en and n in default_idx)
        # embedding: year only - dates, cloud, composite disabled
        self.sp_year.setEnabled(emb)
        self.de_start.setEnabled(not emb)
        self.de_end.setEnabled(not emb)
        self.sp_cloud.setEnabled(opt)
        self.cb_comp.setEnabled(not emb)
        for cb in (self.chk_pct, self.chk_std, self.chk_ndvit):
            cb.setEnabled(opt)
            if not opt:
                cb.setChecked(False)
        self._rebuild_base_checks()

    # ------------------------------------------------------------ AOI UI
    def _aoi_mode_changed(self, idx):
        admin = idx == AOI_ADMIN
        for w in (self.cb_prov, self.cb_dist, self.cb_comm):
            w.setEnabled(admin)
        self.cb_layer.setEnabled(idx == AOI_LAYER)
        self.chk_selected.setEnabled(idx == AOI_LAYER)
        for w in (self.cb_asset, self.btn_afields, self.cb_afield,
                  self.cb_avalue, self.chk_whole_asset):
            w.setEnabled(idx == AOI_ASSET)
        self.btn_draw.setEnabled(idx == AOI_DRAW)
        if admin and self.cb_prov.count() == 0:
            self._populate_provinces()

    # ---- bundled cascading admin units (offline)
    def _populate_provinces(self):
        try:
            from .. import adminmap
            units = adminmap.read_units("province")
        except Exception as e:
            self.lbl_aoi.setText("Bundled boundaries unavailable: {}".format(e))
            return
        if not units:
            self.lbl_aoi.setText(
                "No bundled boundaries found. Copy 'data/boundaries' from "
                "the plugin folder (data/boundaries).")
            return
        self.cb_prov.blockSignals(True)
        self.cb_prov.clear()
        self.cb_prov.addItem("(choose province)", None)
        for label, code, name in units:
            self.cb_prov.addItem(label, (code, name))
        self.cb_prov.blockSignals(False)
        self._cascade_district()

    def _cascade_district(self):
        from .. import adminmap
        self.cb_dist.blockSignals(True)
        self.cb_dist.clear()
        self.cb_dist.addItem("(whole province)", None)
        d = self.cb_prov.currentData()
        if d:
            for label, code, name in adminmap.read_units("district", d[0]):
                self.cb_dist.addItem(label, (code, name))
        self.cb_dist.blockSignals(False)
        self._cascade_commune()

    def _cascade_commune(self):
        from .. import adminmap
        self.cb_comm.blockSignals(True)
        self.cb_comm.clear()
        self.cb_comm.addItem("(whole district)", None)
        d = self.cb_dist.currentData()
        if d:
            for label, code, name in adminmap.read_units("commune", d[0]):
                self.cb_comm.addItem(label, (code, name))
        self.cb_comm.blockSignals(False)

    def _hint_from_geometry(self, g84):
        """Detect which bundled admin unit contains this AOI's centre, so
        the map layout tool can preselect the same area."""
        try:
            from .. import adminmap
            c = g84.centroid().asPoint()
            hint = adminmap.locate_unit(c)
            return hint or {}
        except (AttributeError, TypeError, RuntimeError):  # QGIS geometry
            return {}

    def _admin_selection(self):
        comm = self.cb_comm.currentData()
        dist = self.cb_dist.currentData()
        prov = self.cb_prov.currentData()
        if comm:
            return "commune", comm[0], comm[1]
        if dist:
            return "district", dist[0], dist[1]
        if prov:
            return "province", prov[0], prov[1]
        return None, None, None

    def preview_aoi(self):
        """Resolve the AOI and draw its boundary, without building a stack."""
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            ee = gu.init_ee()
            aoi, name = self.get_aoi(ee)
            if getattr(self, "_local_g84", None) is not None:
                lyr = gu.local_geometry_to_layer(self._local_g84, name)
            else:
                lyr = gu.ee_geometry_to_layer(aoi, name)
            if lyr is not None:
                self._zoom_to_layer(lyr)
            km2 = gu._retry(lambda: aoi.area(1).getInfo()) / 1e6
            self.lbl_aoi.setText(
                "AOI: {}  (~{:,.0f} km2)".format(name, km2))
        except Exception as e:
            self.lbl_aoi.setText("AOI error: {}".format(e))
        finally:
            QApplication.restoreOverrideCursor()

    def _load_asset_fields(self):
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            gu.init_ee()
            self.cb_afield.clear()
            self.cb_afield.addItems(
                gu.asset_field_names(self.cb_asset.currentText().strip()))
        except Exception as e:
            QMessageBox.warning(self, "GEE asset", str(e))
        finally:
            QApplication.restoreOverrideCursor()

    def _load_asset_values(self, field):
        if not field:
            return
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            self.cb_avalue.clear()
            path = self.cb_asset.currentText().strip()
            rows, code_field = gu.asset_value_rows(path, field)
            self._asset_code_field = code_field
            for label, name, code in rows:
                # store (name, code) so we can filter by the UNIQUE code
                self.cb_avalue.addItem(label, (name, code))
            len(rows) - len({r[0] for r in rows})
            if code_field:
                self.lbl_status.setText(
                    "Loaded {} values; duplicate names are shown with their "
                    "parent and code [{}], and filtered by code so the AOI "
                    "is a single unit.".format(len(rows), code_field))
        except Exception as e:
            QMessageBox.warning(self, "GEE asset", str(e))
        finally:
            QApplication.restoreOverrideCursor()

    def _start_draw(self):
        canvas = self.iface.mapCanvas()
        self.poly_tool = PolygonDrawTool(canvas)
        self.poly_tool.finished.connect(self._poly_done)
        canvas.setMapTool(self.poly_tool)
        self.showMinimized()

    def _poly_done(self, geom):
        if geom is None or geom.isEmpty():
            self.showNormal(); self.raise_()
            return
        self.drawn_geom = geom
        try:
            n = len(geom.asPolygon()[0]) - 1
        except (AttributeError, TypeError, RuntimeError):  # Safe fallback
            n = 0
        self.lbl_rect.setText("Polygon set ({} vertices).".format(n))
        self.showNormal(); self.raise_()

    # ------------------------------------------------------------ AOI geometry
    def get_aoi(self, ee):
        """Resolve the AOI for every mode. Sets self._local_g84 when the
        geometry comes from local data (so no GEE call is needed to draw it)."""
        self._local_g84 = None
        idx = self.cb_aoi.currentIndex()

        # -- current canvas extent
        if idx == AOI_CANVAS:
            canvas = self.iface.mapCanvas()
            ext = canvas.extent()
            if ext.isEmpty():
                raise RuntimeError("The map canvas has no extent yet.")
            crs = canvas.mapSettings().destinationCrs()
            g = QgsGeometry.fromRect(ext)
            self._local_g84 = gu.to_wgs84_geometry(g, crs)
            hint = self._hint_from_geometry(self._local_g84)
            self.state["admin_hint"] = hint
            label = "Canvas extent"
            if hint.get("name"):
                label += " ({})".format(hint["name"])
            return gu.qgis_geom_to_ee(self._local_g84, gu.WGS84), label

        # -- bundled administrative unit
        if idx == AOI_ADMIN:
            from .. import adminmap
            level, code, name = self._admin_selection()
            if level is None:
                raise RuntimeError("Choose at least a province.")
            unit, nf, cf = adminmap._unit_layer(level, code)
            if unit is None or unit.featureCount() == 0:
                raise RuntimeError("Could not load the bundled boundary.")
            feats = list(unit.getFeatures())
            g = QgsGeometry(feats[0].geometry())
            for ft in feats[1:]:
                g = g.combine(ft.geometry())
            self._local_g84 = gu.to_wgs84_geometry(g, unit.crs())
            self.state["admin_hint"] = {
                "level": level, "name": name, "code": str(code)}
            return gu.qgis_geom_to_ee(self._local_g84, gu.WGS84), name

        # -- active layer, selected features
        if idx == AOI_ACTIVE:
            from qgis.core import QgsVectorLayer
            lyr = self.iface.activeLayer()
            if not isinstance(lyr, QgsVectorLayer):
                raise RuntimeError(
                    "Select a vector layer as the active layer.")
            feats = list(lyr.selectedFeatures())
            if not feats:
                raise RuntimeError(
                    "No features selected in the active layer.")
            g = QgsGeometry(feats[0].geometry())
            for ft in feats[1:]:
                g = g.combine(ft.geometry())
            self._local_g84 = gu.to_wgs84_geometry(g, lyr.crs())
            self.state["admin_hint"] = {}
            return gu.qgis_geom_to_ee(self._local_g84, gu.WGS84), lyr.name()

        # -- a chosen polygon layer
        if idx == AOI_LAYER:
            lyr = self.cb_layer.currentLayer()
            if lyr is None:
                raise RuntimeError("Please choose a polygon layer.")
            feats = (lyr.selectedFeatures() if self.chk_selected.isChecked()
                     and lyr.selectedFeatureCount() > 0
                     else list(lyr.getFeatures()))
            if not feats:
                raise RuntimeError("Layer has no features / no selection.")
            g = QgsGeometry(feats[0].geometry())
            for ft in feats[1:]:
                g = g.combine(ft.geometry())
            self._local_g84 = gu.to_wgs84_geometry(g, lyr.crs())
            self.state["admin_hint"] = {}
            return gu.qgis_geom_to_ee(self._local_g84, gu.WGS84), lyr.name()

        # -- GEE asset (advanced)
        if idx == AOI_ASSET:
            path = self.cb_asset.currentText().strip()
            fc = gu.asset_fc(path)
            if self.chk_whole_asset.isChecked():
                return fc.geometry(), os.path.basename(path)
            field = self.cb_afield.currentText()
            data = self.cb_avalue.currentData()
            typed = self.cb_avalue.currentText().strip()
            if not field or (data is None and not typed):
                raise RuntimeError(
                    "Click 'Load fields', then choose field and value "
                    "(or tick 'whole asset').")
            code_field = getattr(self, "_asset_code_field", None)
            name, code = data if data is not None else (typed, "")

            def _nonempty(cand):
                try:
                    return gu._retry(lambda: cand.size().getInfo()) > 0
                except (ImportError, ModuleNotFoundError, RuntimeError):  # GEE/Import
                    return False

            filt = None
            if code_field and code:
                for cv in (_coerce(code), str(code)):
                    cand = fc.filter(ee.Filter.eq(code_field, cv))
                    if _nonempty(cand):
                        filt = cand
                        break
            if filt is None:
                cand = fc.filter(ee.Filter.eq(field, name))
                if _nonempty(cand):
                    filt = cand
            if filt is None:
                raise RuntimeError(
                    "No feature matched (field '{}', value '{}').".format(
                        field, name))
            base = os.path.basename(path).lower()
            level = ("province" if "khet" in base else
                     "district" if "srok" in base else
                     "commune" if "khum" in base else None)
            self.state["admin_hint"] = {
                "level": level, "name": name, "code": str(code or "")}
            geom = filt.geometry()
            if gu._retry(lambda: geom.area(1).getInfo()) <= 0:
                raise RuntimeError("The selected unit has empty geometry.")
            return geom, ("{} [{}]".format(name, code) if code else name)

        # -- drawn polygon
        if self.drawn_geom is None:
            raise RuntimeError("Please draw a polygon first.")
        canvas_crs = self.iface.mapCanvas().mapSettings().destinationCrs()
        self._local_g84 = gu.to_wgs84_geometry(self.drawn_geom, canvas_crs)
        self.state["admin_hint"] = self._hint_from_geometry(self._local_g84)
        return (gu.qgis_geom_to_ee(self._local_g84, gu.WGS84),
                "Drawn polygon")

    # ------------------------------------------------------------ stack build
    def build_stack(self):
        ee = gu.init_ee()
        aoi, aoi_name = self.get_aoi(ee)
        start = self.de_start.date().toString("yyyy-MM-dd")
        end = self.de_end.date().toString("yyyy-MM-dd")
        cloud = self.sp_cloud.value()
        method = self.cb_comp.currentText()
        i = self.cb_sensor.currentIndex()
        sel_opt = [n for n in gu.OPTICAL_INDICES
                   if self.idx_checks[n].isChecked()]
        sel_re = [n for n in gu.S2_ONLY_INDICES
                  if self.idx_checks[n].isChecked()]
        sel_sar = [n for n in gu.SAR_INDICES
                   if self.idx_checks[n].isChecked()]

        stack, scale, kind = None, 10, "optical"
        if i in (SEN_S2, SEN_L8, SEN_L9, SEN_L89):
            code = {SEN_S2: "S2", SEN_L8: "L8",
                    SEN_L9: "L9", SEN_L89: "L89"}[i]
            base = gu.S2_BASE if code == "S2" else gu.LS_BASE
            col, scale = gu.optical_collection(code, aoi, start, end, cloud)
            img = gu.composite(col, method)
            names = sel_opt + (sel_re if code == "S2" else [])
            stack = gu.add_optical_indices(img, names, code)
            stack = self._add_temporal(stack, col, base)
        elif i == SEN_S1:
            col, scale = gu.s1_collection(aoi, start, end)
            stack = gu.add_sar_indices(gu.composite(col, method), sel_sar)
            kind = "sar"
        elif i == SEN_FUSION:
            col2, scale = gu.optical_collection("S2", aoi, start, end, cloud)
            opt = gu.add_optical_indices(
                gu.composite(col2, method), sel_opt + sel_re, "S2")
            opt = self._add_temporal(opt, col2, gu.S2_BASE)
            col1, _ = gu.s1_collection(aoi, start, end)
            sar = gu.add_sar_indices(gu.composite(col1, method), sel_sar)
            stack = opt.addBands(sar)
            kind = "fusion"
        else:                                          # Satellite Embedding
            year = self.sp_year.value()
            img, size = gu.embedding_image(aoi, year)
            if size.getInfo() == 0:
                raise RuntimeError(
                    "No Satellite Embedding image for {} here. The dataset "
                    "starts in 2017 and the newest year may not be published "
                    "yet - choose an earlier year.".format(year))
            stack, scale, kind = img, 10, "embedding"
            aoi_name = "{} (embedding {})".format(aoi_name, year)

        if self.chk_lst.isChecked():
            stack = stack.addBands(gu.landsat_lst(aoi, start, end))
        if self.chk_terrain.isChecked():
            stack = stack.addBands(gu.terrain_bands())

        if self.chk_important.isChecked() and kind != "embedding":
            stack = gu.select_important_bands(stack)

        # trim unchecked ORIGINAL bands from the OUTPUT only. Indices are
        # already computed above, so trimming here is purely cosmetic for
        # the output band list and never breaks index math.
        if kind != "embedding" and self.base_checks:
            drop = [n for n, cb in self.base_checks.items()
                    if not cb.isChecked()]
            if drop:
                import ee
                names = stack.bandNames()
                keep_names = names.filter(
                    ee.Filter.inList("item", ee.List(drop)).Not())
                stack = stack.select(keep_names)

        stack = stack.clip(aoi).toFloat()
        # keep the local WGS84 AOI so the Map Layout tool can reuse it
        self.state["aoi_geom_wgs84"] = (
            QgsGeometry(self._local_g84)
            if getattr(self, "_local_g84", None) is not None else None)
        self.state.update({"stack": stack, "aoi": aoi, "aoi_name": aoi_name,
                           "scale": scale, "kind": kind})
        return stack, aoi, aoi_name, scale, kind, sel_opt + sel_re + sel_sar

    def _add_temporal(self, stack, col, base_bands):
        if self.chk_pct.isChecked():
            stack = gu.add_percentile_bands(stack, col, base_bands)
        if self.chk_std.isChecked():
            stack = gu.add_std_bands(stack, col, base_bands)
        if self.chk_ndvit.isChecked():
            stack = gu.add_temporal_ndvi(stack, col)
        return stack

    # ------------------------------------------------------------ run
    def run(self):
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            self.lbl_status.setText("Building stack on GEE ...")
            QApplication.processEvents()
            stack, aoi, aoi_name, scale, kind, indices = self.build_stack()
            band_names = gu._retry(lambda: stack.bandNames().getInfo())
            if not band_names:
                raise RuntimeError(
                    "The stack has no bands - usually the AOI is empty or "
                    "no scenes matched the date/cloud filter. Check the AOI "
                    "value, widen the date range, or raise the cloud %.")
            self.state["band_names"] = band_names
            gamma = self.GAMMA

            if kind == "sar":
                vis = {"bands": ["VV", "VH", "VV"], "min": -25, "max": 0}
            elif kind == "embedding":
                vis = {"bands": ["A01", "A16", "A09"],
                       "min": -0.3, "max": 0.3}
            else:
                vis = {"bands": ["red", "green", "blue"],
                       "min": 0.02, "max": 0.35, "gamma": gamma}
            gu.add_ee_image_layer(
                stack, vis, "{} - {}".format(
                    self.cb_sensor.currentText().split(" (")[0], aoi_name))

            if self.chk_add_idx_layers.isChecked() and kind != "embedding":
                for n in indices:
                    if n in band_names:
                        gu.add_ee_image_layer(
                            stack.select(n),
                            {"min": -1, "max": 1,
                             "palette": ["8c510a", "f5f5f5", "01665e"]},
                            "Index: " + n)

            gu.raise_points_above(("sample_points",))
            aoi_note = ""
            try:
                if getattr(self, "_local_g84", None) is not None:
                    # local AOI (shapefile / rectangle): no GEE call needed
                    aoi_lyr = gu.local_geometry_to_layer(
                        self._local_g84, aoi_name)
                else:
                    aoi_lyr = gu.ee_geometry_to_layer(aoi, aoi_name)
                if aoi_lyr is None:
                    raise RuntimeError("layer invalid")
                self._zoom_to_layer(aoi_lyr)
            except Exception as e:
                aoi_note = ("\nAOI boundary layer could not be built "
                            "({}). You can still continue: use your asset "
                            "boundary or shapefile directly as the AOI "
                            "polygon in Tool 2.".format(e))

            self.lbl_status.setText(
                "Done. Stack has {} bands - preview tiles on the canvas, "
                "stored in memory for Tool 4.{}".format(
                    len(band_names), aoi_note))
            if self.chk_local.isChecked():
                self._start_local_download(
                    stack, aoi, aoi_name, scale, kind, band_names, gamma)
        except Exception as e:
            self.lbl_status.setText("Error: {}".format(e))
            QMessageBox.critical(self, "Error",
                                 "{}\n\n{}".format(e, traceback.format_exc()))
        finally:
            QApplication.restoreOverrideCursor()

    def _start_local_download(self, stack, aoi, aoi_name, scale, kind,
                              band_names, gamma):
        """Download the local multiband GeoTIFF in a BACKGROUND task so
        QGIS stays responsive (the preview tiles are already loaded)."""
        safe = re.sub(r"[^A-Za-z0-9_-]+", "_", aoi_name)[:40]
        path = os.path.join(tempfile.gettempdir(),
                            "gee_stack_{}.tif".format(safe))
        dlg = self

        class _DlTask(QgsTask):
            def __init__(t):
                super().__init__("GEE: local stack download ({})"
                                 .format(aoi_name), QgsTask.Flag.CanCancel)
                t.error = None
                t.n_tiles = 0
                t.scale_use = scale
                t.msg = "starting"

            def run(t):
                try:
                    def prog(frac, msg):
                        t.msg = msg
                        t.setProgress(min(99.0, frac * 100.0))
                        if t.isCanceled():
                            raise RuntimeError("canceled by user")
                    t.n_tiles, t.scale_use = gu.download_stack_smart(
                        stack, aoi, scale, path, len(band_names),
                        progress=prog)
                    return True
                except Exception as e:
                    t.error = str(e)
                    return False

            def finished(t, ok):
                dlg._local_done(ok, t, path, aoi_name, scale, kind,
                                band_names, gamma)

        self._dl_task = _DlTask()
        self._dl_task.progressChanged.connect(
            lambda p: self.lbl_status.setText(
                "Local stack: {}  ({} %)  -  QGIS stays usable; preview "
                "tiles are already on the canvas.".format(
                    self._dl_task.msg, int(p))))
        QgsApplication.taskManager().addTask(self._dl_task)
        self.lbl_status.setText(
            "Local stack download started in the background - QGIS stays "
            "usable. Preview tiles are already on the canvas.")

    def _local_done(self, ok, task, path, aoi_name, scale, kind,
                    band_names, gamma):
        """Runs on the main thread when the background download ends."""
        if not ok:
            self.lbl_status.setText(
                "Local GeoTIFF failed ({}). The preview tiles and the "
                "in-memory stack for Tool 4 are unaffected - you can "
                "click Build stack again to retry, or use Export to "
                "Drive.".format(task.error))
            return
        try:
            with open(path + ".bands.txt", "w", encoding="utf-8") as fh:
                for i, b in enumerate(band_names, start=1):
                    fh.write("Band {}: {}\n".format(i, b))
            gu.set_gdal_band_descriptions(path, band_names)
            gu.write_band_qml(path, band_names, kind, gamma)
            lyr = QgsRasterLayer(path, "Stack {} bands - {}".format(
                len(band_names), aoi_name))
            if not lyr.isValid():
                self.lbl_status.setText("Local GeoTIFF could not be loaded.")
                return
            QgsProject.instance().addMapLayer(lyr)
            gu.set_raster_band_names(lyr, band_names)
            gu.apply_rgb_style(lyr, band_names, kind, gamma)
            note = ""
            if task.scale_use > scale:
                note += (" - local copy resampled to {} m; the in-memory "
                         "stack and Drive export stay at {} m"
                         .format(task.scale_use, scale))
            if task.n_tiles > 1:
                note += " ({} tiles merged)".format(task.n_tiles)
            self.lbl_status.setText(
                "Local multiband stack loaded at {} m{}: {}\n"
                "Change R/G/B bands: Layer Properties > Symbology > "
                "Multiband color. Band list: {}.bands.txt".format(
                    task.scale_use, note, path, path))
        except Exception as e:
            self.lbl_status.setText(
                "Local GeoTIFF failed while loading: {}".format(e))

    # ------------------------------------------------------------ Drive export
    def export_drive(self):
        if self.state.get("stack") is None:
            QMessageBox.information(self, "No stack",
                                    "Run 'Build stack' first.")
            return
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            name = gu.export_to_drive(
                self.state["stack"], self.state["aoi"], self.state["scale"],
                "gee_stack_" + (self.state.get("aoi_name") or "aoi"))
            self.lbl_status.setText(
                "Export task '{}' started (full resolution, {} m).\n"
                "The GeoTIFF will appear in your Google Drive folder "
                "'{}' in some minutes. You can watch progress in the "
                "Tasks tab of the GEE Code Editor.".format(
                    name, self.state["scale"], gu.DRIVE_FOLDER))
        except Exception as e:
            QMessageBox.critical(self, "Drive export failed", str(e))
        finally:
            QApplication.restoreOverrideCursor()

    def _zoom_to_layer(self, lyr):
        """Zoom the canvas to the AOI so the boundary is always visible."""
        try:
            canvas = self.iface.mapCanvas()
            ct = QgsCoordinateTransform(
                lyr.crs(), canvas.mapSettings().destinationCrs(),
                QgsProject.instance())
            ext = ct.transformBoundingBox(lyr.extent())
            ext.scale(1.08)
            canvas.setExtent(ext)
            canvas.refresh()
        except (AttributeError, TypeError):  # Qt/QGIS API compatibility
            pass

    # ------------------------------------------------------------ Asset
    ASSET_FOLDER_KEY = "rupp_gee_lulc/asset_folder"

    def export_asset(self):
        if self.state.get("stack") is None:
            QMessageBox.information(self, "No stack",
                                    "Run 'Build stack' first.")
            return
        folder = QSettings().value(
            self.ASSET_FOLDER_KEY, "users/yamsarath168/QGIS_stacks",
            type=str)
        safe = re.sub(r"[^A-Za-z0-9_-]+", "_",
                      self.state.get("aoi_name") or "aoi")[:40]
        default = "{}/gee_stack_{}".format(folder.rstrip("/"), safe)
        asset_id, ok = QInputDialog.getText(
            self, "Export stack to GEE Asset",
            "Asset ID (the parent folder must exist in your GEE account):",
            text=default)
        if not ok or not asset_id.strip():
            return
        asset_id = asset_id.strip()
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            name = gu.export_to_asset(
                self.state["stack"], self.state["aoi"],
                self.state["scale"], asset_id)
            QSettings().setValue(self.ASSET_FOLDER_KEY,
                                 asset_id.rsplit("/", 1)[0])
            self.state["last_asset_id"] = asset_id
            self.lbl_status.setText(
                "Asset export task '{}' started at {} m (full "
                "resolution).\nWatch progress in the Tasks tab of the GEE "
                "Code Editor. When it finishes, click 'Load stack from "
                "Asset ...' to bring it back - even in a new QGIS session, "
                "no rebuild needed.".format(name, self.state["scale"]))
        except Exception as e:
            QMessageBox.critical(self, "Asset export failed", str(e))
        finally:
            QApplication.restoreOverrideCursor()

    def load_asset(self):
        default = self.state.get("last_asset_id") or QSettings().value(
            self.ASSET_FOLDER_KEY, "users/yamsarath168/QGIS_stacks",
            type=str)
        asset_id, ok = QInputDialog.getText(
            self, "Load stack from GEE Asset", "Asset ID:", text=default)
        if not ok or not asset_id.strip():
            return
        asset_id = asset_id.strip()
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            self.lbl_status.setText("Loading asset ...")
            QApplication.processEvents()
            gu.init_ee()
            img, band_names, scale, kind = gu.load_image_asset(asset_id)
            aoi = img.geometry()
            aoi_name = asset_id.rsplit("/", 1)[-1]
            img = img.toFloat()
            self.state.update({
                "stack": img, "aoi": aoi, "aoi_name": aoi_name,
                "scale": scale, "kind": kind, "band_names": band_names,
                "last_asset_id": asset_id})
            gamma = self.GAMMA
            if kind == "sar":
                vis = {"bands": ["VV", "VH", "VV"], "min": -25, "max": 0}
            elif kind == "embedding":
                vis = {"bands": ["A01", "A16", "A09"],
                       "min": -0.3, "max": 0.3}
            else:
                vis = {"bands": ["red", "green", "blue"],
                       "min": 0.02, "max": 0.35, "gamma": gamma}
            gu.add_ee_image_layer(img, vis, "Asset - " + aoi_name)
            self.lbl_status.setText(
                "Asset loaded: {} bands at {} m. Ready for Tool 4 - "
                "no rebuild needed.".format(len(band_names), scale))
            if self.chk_local.isChecked():
                self._start_local_download(
                    img, aoi, aoi_name, scale, kind, band_names, gamma)
        except Exception as e:
            self.lbl_status.setText("Error: {}".format(e))
            QMessageBox.critical(self, "Load asset failed", str(e))
        finally:
            QApplication.restoreOverrideCursor()

    # ------------------------------------------------------------ reset
    def reset_defaults(self):
        self.cb_sensor.setCurrentIndex(SEN_S2)
        self.de_start.setDate(QDate.currentDate().addYears(-1))
        self.de_end.setDate(QDate.currentDate())
        self.sp_year.setValue(QDate.currentDate().year() - 1)
        self.sp_cloud.setValue(30)
        self.cb_comp.setCurrentIndex(0)
        self.cb_aoi.setCurrentIndex(AOI_ADMIN)
        self.chk_whole_asset.setChecked(False)
        self.chk_selected.setChecked(True)
        self.drawn_geom = None
        self.lbl_rect.setText("No polygon drawn.")
        for cb in (self.chk_pct, self.chk_std, self.chk_ndvit,
                   self.chk_lst, self.chk_terrain):
            cb.setChecked(False)
        self.chk_important.setChecked(False)
        self.chk_local.setChecked(True)
        self.chk_add_idx_layers.setChecked(False)
        self._sensor_changed()
        self.lbl_status.setText("Reset to defaults.")
