# -*- coding: utf-8 -*-
"""Tool 5: LULC Admin Map.

Uses the bundled Admin Map engine (adminmap.py + layout_builder.py) and:
  1. inserts the classified LULC layer into the map (under boundaries,
     above the basemap),
  2. makes sure the title says LULC,
  3. adds the accuracy assessment (OA, Kappa, PA, UA) in a box below
     the legend,
  4. adds an LULC class legend that works even when the classification
     is a GEE tile layer.

UPDATED: Simplified - no "Layers to include", always uses surround mode.
"""

from qgis.PyQt.QtWidgets import QInputDialog, QMessageBox
from qgis.core import (
    QgsProject, QgsRasterLayer, QgsVectorLayer, QgsFillSymbol,
    QgsCategorizedSymbolRenderer, QgsRendererCategory,
    QgsPalettedRasterRenderer,
)

from . import gee_utils as gu


def _auto_classified_layer(iface, state):
    """Find the classified layer without prompting, in priority order:
    1. the stored local layer id (set once Tool 4's download finishes),
    2. any raster layer whose name starts with 'LULC ' (the local map or
       the GEE preview tiles Tool 4 added synchronously),
    3. the most recently added LULC-like raster.
    Returns None if nothing obvious is found (then the picker is used)."""
    proj = QgsProject.instance()

    # First check: stored layer ID from state
    lid = state.get("classified_layer_id")
    if lid:
        lyr = proj.mapLayer(lid)
        if lyr is not None:
            return lyr

    # Second check: any raster with LULC in name
    rasters = [lyr for lyr in proj.mapLayers().values()
               if isinstance(lyr, QgsRasterLayer)]

    def _score(lyr):
        n = lyr.name().lower()
        if "(local)" in n and n.startswith("lulc"):
            return 0
        if n.startswith("lulc "):
            return 1
        if "lulc" in n or "class" in n or "classification" in n:
            return 2
        if "gee" in n and ("tile" in n or "preview" in n):
            return 3
        return 9

    cand = sorted(rasters, key=_score)
    if cand and _score(cand[0]) <= 3:
        return cand[0]
    return None


def _pick_classified_layer(iface):
    """Let the user pick the classified layer (GeoTIFF or GEE tiles)."""
    rasters = [lyr for lyr in QgsProject.instance().mapLayers().values()
               if isinstance(lyr, QgsRasterLayer)]
    if not rasters:
        raise RuntimeError("No raster layers in the project. Run Tool 4 "
                           "first (or load the classified GeoTIFF).")
    # put likely candidates first
    rasters.sort(key=lambda r: 0 if "lulc" in r.name().lower()
                 or "class" in r.name().lower() else 1)
    names = [lyr.name() for lyr in rasters]
    name, ok = QInputDialog.getItem(
        iface.mainWindow(), "LULC Map Layout",
        "Classified LULC layer:", names, 0, False)
    if not ok:
        return None
    return rasters[names.index(name)]


def _lulc_classes_for_legend(scheme="lulc"):
    """Classes for the legend, following the active scheme (LULC / WorldCover
    / custom). For the LULC scheme, Water is recoloured to match the Water
    bodies vector (#96c8e6) so map and legend agree."""
    out = []
    for code, name, color in gu.scheme_classes(scheme):
        if scheme == "lulc" and name.lower().startswith("water"):
            color = WATER_VECTOR_COLOR
        out.append((code, name, color))
    return out


def _legend_proxy_layer(scheme="lulc"):
    """A small categorized polygon layer so the classes appear in the legend
    even when the classification is a GEE tile (WMS) layer."""
    mem = QgsVectorLayer("Polygon?crs=EPSG:4326&field=class_code:integer",
                         "LULC Classes", "memory")
    cats = []
    for code, name, color in _lulc_classes_for_legend(scheme):
        sym = QgsFillSymbol.createSimple(
            {"color": color, "outline_color": "black",
             "outline_width": "0.15"})
        cats.append(QgsRendererCategory(code, sym, name))
    mem.setRenderer(QgsCategorizedSymbolRenderer("class_code", cats))
    QgsProject.instance().addMapLayer(mem, False)
    return mem


def _split_admin_code(code):
    """Cambodia admin codes are hierarchical:
    province = 2 digits, district = 4, commune = 6.
    Returns (prov_code, dist_code, comm_code) as far as available."""
    c = "".join(ch for ch in str(code) if ch.isdigit())
    prov = c[:2] if len(c) >= 2 else ""
    dist = c[:4] if len(c) >= 4 else ""
    comm = c[:6] if len(c) >= 6 else ""
    return prov, dist, comm


def _find_by_code(cmb, code):
    """Index of the combo item whose stored code == code."""
    if not code:
        return -1
    for i in range(cmb.count()):
        d = cmb.itemData(i)
        if d and str(d[0]) == str(code):
            return i
    for i in range(cmb.count()):
        d = cmb.itemData(i)
        if d and str(d[0]).lstrip("0") == str(code).lstrip("0"):
            return i
    return -1


def _find_by_name(cmb, name):
    up = (name or "").strip().upper()
    if not up:
        return -1
    for i in range(cmb.count()):
        d = cmb.itemData(i)
        if d and str(d[1]).strip().upper() == up:
            return i
    for i in range(cmb.count()):
        d = cmb.itemData(i)
        if d and up in str(d[1]).strip().upper():
            return i
    return -1


def _preselect_area(dlg, hint):
    """Preselect province/district/commune from the AOI chosen in Tool 1."""
    name = (hint.get("name") or "").strip()
    level = hint.get("level")
    code = hint.get("code") or ""
    prov_c, dist_c, comm_c = _split_admin_code(code)
    resolved = {"prov": "", "dist": "", "comm": ""}
    if not level:
        return resolved

    # Map level to admin level index
    level_map = {"province": 0, "district": 1, "commune": 2,
                 "admin0": 0, "admin1": 1, "admin2": 2, "admin3": 3}
    lvl_idx = level_map.get(level, 0)

    # Level 0 / province
    ip = _find_by_code(dlg.cmb_prov, prov_c) if prov_c else -1
    if ip < 0 and lvl_idx == 0:
        ip = _find_by_name(dlg.cmb_prov, name)
    if ip > 0:
        dlg.cmb_prov.setCurrentIndex(ip)
        d = dlg.cmb_prov.currentData()
        resolved["prov"] = d[1] if d else ""
    if lvl_idx == 0:
        return resolved

    # Level 1 / district
    idd = _find_by_code(dlg.cmb_dist, dist_c) if dist_c else -1
    if idd < 0 and lvl_idx == 1:
        idd = _find_by_name(dlg.cmb_dist, name)
    if idd > 0:
        dlg.cmb_dist.setCurrentIndex(idd)
        d = dlg.cmb_dist.currentData()
        resolved["dist"] = d[1] if d else ""
    if lvl_idx == 1:
        return resolved

    # Level 2 / commune
    ic = _find_by_code(dlg.cmb_comm, comm_c) if comm_c else -1
    if ic < 0:
        ic = _find_by_name(dlg.cmb_comm, name)
    if ic > 0:
        dlg.cmb_comm.setCurrentIndex(ic)
        d = dlg.cmb_comm.currentData()
        resolved["comm"] = d[1] if d else ""
    return resolved


# colors requested for the LULC Admin Map insets
INSET1_2ND_COLOR = "#000000"
INSET2_2ND_COLOR = "#000000"
INSET1_2ND_FILL = "BDiagonal"
INSET2_2ND_FILL = "Diagonal X"
WATER_VECTOR_COLOR = "#96c8e6"


def _auto_configure_insets(dlg, resolved):
    """Set up the two insets to mirror the selected area."""
    prov = resolved.get("prov") or ""
    dist = resolved.get("dist") or ""
    comm = resolved.get("comm") or ""

    if prov:
        dlg.configure_inset(
            0, layer_name="KH_Provinces", field="Name_EN", value=prov,
            color="#ffe091",
            layer2_name=("KH_Districts" if dist else None),
            field2=("Name_EN" if dist else None),
            value2=(dist if dist else None),
            color2=INSET1_2ND_COLOR, fill_style2=INSET1_2ND_FILL,
            labels="highlight", zoom="whole", clip_focus=False,
            enabled=True)

    if dist and prov:
        dlg.configure_inset(
            1, layer_name="KH_Districts", field="Name_EN", value=dist,
            color="#ffe091",
            layer2_name=("KH_Communes" if comm else None),
            field2=("Name_EN" if comm else None),
            value2=(comm if comm else None),
            color2=INSET2_2ND_COLOR, fill_style2=INSET2_2ND_FILL,
            labels="highlight", zoom="whole", clip_focus=False,
            clip_to_parent=prov, enabled=True)


def _remove_previous_lulc_layers():
    """Remove layers this tool added in a previous run."""
    proj = QgsProject.instance()
    for lyr in list(proj.mapLayers().values()):
        nm = lyr.name()
        if (nm.startswith("LULC Classes")
                or nm.startswith("AOI - ")
                or nm.endswith("(local)") and nm.startswith("LULC ")):
            proj.removeMapLayer(lyr.id())


def run_lulc_adminmap(iface, state):
    from . import adminmap
    from .adminmap_dialog import AdminMapDialog

    _remove_previous_lulc_layers()
    if not adminmap.read_units(adminmap.ADMIN_LEVELS[0]):
        raise RuntimeError(
            "Bundled boundary shapefiles were not found.\n"
            "Copy the folders 'data/boundaries', 'data/context', "
            "'data/styles' and 'assets' from your Admin Map plugin into "
            "the rupp_gee_lulc plugin folder, then try again.")

    classified = None
    lid = state.get("classified_layer_id")
    if lid:
        classified = QgsProject.instance().mapLayer(lid)
    if classified is None:
        classified = _pick_classified_layer(iface)
    if classified is None:
        return None

    # local GeoTIFF: make sure it uses the LULC palette
    is_gdal = classified.providerType() == "gdal"
    scheme = state.get("scheme") or "lulc"
    if is_gdal and not isinstance(classified.renderer(),
                                  QgsPalettedRasterRenderer):
        gu.apply_lulc_style(classified, scheme)
    legend_proxy = None if is_gdal else _legend_proxy_layer(scheme)

    dlg = AdminMapDialog(iface, iface.mainWindow())
    prepare_lulc_dialog(iface, state, dlg)
    if not dlg.exec():
        return None
    return generate_lulc_map(iface, state, dlg, classified,
                             legend_proxy, is_gdal)


def prepare_lulc_dialog(iface, state, dlg):
    """Preselect the area / insets on an AdminMapDialog from session state."""
    try:
        from qgis.PyQt.QtWidgets import QApplication
        from qgis.PyQt.QtCore import Qt
        QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
        hint = state.get("admin_hint") or {}
        resolved = _preselect_area(dlg, hint)
        _auto_configure_insets(dlg, resolved)
        aoi_name = state.get("aoi_name") or ""
        is_custom_aoi = (aoi_name.startswith("Canvas extent")
                         or aoi_name.startswith("Drawn"))
        g84 = state.get("aoi_geom_wgs84")
        if g84 is not None and is_custom_aoi and not hint.get("level"):
            dlg.set_custom_geometry(g84, aoi_name or "AOI")
        else:
            try:
                dlg.cmb_source.setCurrentIndex(0)
            except (AttributeError, TypeError):  # Qt/QGIS API compatibility
                pass
    except (AttributeError, TypeError):  # Qt/QGIS API compatibility
        pass
    finally:
        try:
            QApplication.restoreOverrideCursor()
        except (AttributeError, TypeError):  # Qt cursor
            pass


def generate_lulc_map(iface, state, dlg, classified=None,
                      legend_proxy=None, is_gdal=None):
    """Generate the LULC map layout from a configured AdminMapDialog."""
    from . import adminmap
    from .layout_builder import build_layout

    if classified is None:
        classified = _auto_classified_layer(iface, state)
        if classified is None:
            classified = _pick_classified_layer(iface)
        if classified is None:
            return None

    scheme = state.get("scheme") or "lulc"
    if is_gdal is None:
        is_gdal = classified.providerType() == "gdal"
        if is_gdal and not isinstance(classified.renderer(),
                                      QgsPalettedRasterRenderer):
            gu.apply_lulc_style(classified, scheme)
    if legend_proxy is None and not is_gdal:
        legend_proxy = _legend_proxy_layer(scheme)

    p = dlg.params()
    if not p.get("level") and p.get("custom_geom") is None:
        QMessageBox.warning(
            iface.mainWindow(), "LULC Map Layout",
            "Choose an administrative unit, or pick 'Current map canvas "
            "extent' / 'Drawn polygon' as the area source.")
        return None

    report = state.get("report_summary") or ""

    def wrapped_build_layout(iface2, lp):
        # Title with LULC + admin-level suffix
        unit = p.get("unit_name") or ""
        level = p.get("level") or ""
        suffix_map = {"commune": " Commune", "district": " District",
                      "province": " Province",
                      "admin3": " Commune", "admin2": " District",
                      "admin1": " Province", "admin0": ""}
        suffix = suffix_map.get(level, "")
        if suffix and unit and not unit.lower().rstrip().endswith(
                suffix.strip().lower()):
            unit_titled = unit + suffix
        else:
            unit_titled = unit
        title = p.get("title") or "LULC Map of {}".format(unit_titled)
        if "lulc" not in title.lower():
            title = "LULC - " + title
        lp["map_title"] = title
        lp["layout_name"] = "LULC Map - " + (unit or "")

        # Classified layer under boundaries/context, above basemap
        ids = list(lp.get("main_layer_ids") or [])
        if p.get("basemap") and ids:
            ids.insert(len(ids) - 1, classified.id())
        else:
            ids.append(classified.id())
        lp["main_layer_ids"] = ids

        # Legend entries for LULC classes
        legend_lyr = legend_proxy if legend_proxy else classified
        root = QgsProject.instance().layerTreeRoot()
        node = root.findLayer(legend_lyr.id())
        if node is not None:
            (node.parent() or root).removeChildNode(node)
        grp = root.findGroup("Main") or root
        pos = None
        for j, child in enumerate(grp.children()):
            try:
                child_lyr = QgsProject.instance().mapLayer(child.layerId())
            except (AttributeError, TypeError, RuntimeError):  # QGIS core
                child_lyr = None
            if child_lyr is not None and \
                    child_lyr.name() == "Water bodies":
                pos = j + 1
        if pos is None:
            grp.addLayer(legend_lyr)
        else:
            grp.insertLayer(pos, legend_lyr)

        lids = list(lp.get("legend_layer_ids") or [])
        bids = set(lp.get("boundary_layer_ids") or [])
        insert_at = 0
        for i, lid in enumerate(lids):
            if lid in bids:
                insert_at = i + 1
        lids.insert(insert_at, legend_lyr.id())
        lp["legend_layer_ids"] = lids

        # Accuracy assessment box
        lp["report_text"] = report
        return build_layout(iface2, lp)

    layout = adminmap.build_admin_map(iface, p, wrapped_build_layout)

    # Show classified layer in layer tree Main group
    try:
        root = QgsProject.instance().layerTreeRoot()
        grp = root.findGroup("Main")
        if grp and root.findLayer(classified.id()) is None:
            grp.addLayer(classified)
    except (ImportError, ModuleNotFoundError, RuntimeError):  # GEE/Import
        pass

    # Open print layout designer
    try:
        if layout is not None:
            iface.openLayoutDesigner(layout)
    except (OSError, IOError):  # File operation
        pass
    return layout
