# Copyright (C) 2026 Wilhelmus Dingenus de Pooter
# SPDX-License-Identifier: GPL-3.0-or-later
"""
enc_s111.py
-----------
S-111 Surface Currents — read + portray a time-step as a current-arrow overlay.

S-111 is an S-100 HDF5 coverage (GDAL native ``S111`` driver): a time SERIES of
regularly-gridded arrays, each Group_NNN holding two bands — surfaceCurrentSpeed
(knots) and surfaceCurrentDirection (degrees, from true north, clockwise). This
module detects S-111 datasets (folder or .zip, like enc_s102), exposes the time
axis, and builds ONE time-step as a point ``QgsVectorLayer`` of current arrows,
portrayed per the official S-111 Portrayal Catalogue: an arrow per grid node,
rotated to the current direction, coloured by a 9-band speed scale.

Rendered as a NATIVE overlay in the Chart Viewer (like S-102 / SOUNDG), driven by
a bottom time slider.
"""

import os
import re
import shutil
import tempfile
import zipfile
from datetime import datetime, timedelta

try:
    from qgis.core import (
        Qgis, QgsMessageLog, QgsVectorLayer, QgsField, QgsFeature, QgsGeometry,
        QgsPointXY, QgsCoordinateReferenceSystem, QgsCoordinateTransform,
        QgsProject, QgsMarkerSymbol, QgsSvgMarkerSymbolLayer, QgsProperty,
        QgsRuleBasedRenderer, QgsSymbolLayer)
    from qgis.PyQt.QtCore import QVariant
except (ImportError, ValueError):                       # headless import guard
    Qgis = None


# Official S-111 day-palette portrayal, straight from the S-111 Portrayal
# Catalogue (Rules/select_arrow.xsl + Symbols/SVGStyle_S111day.css): per speed
# band (knots; lower inclusive, upper exclusive, last open-ended) the arrow fill
# STEP colour AND the SIZE scale — bands 1-3 fixed 0.40×, bands 4-8 scale with
# surfaceCurrentSpeed × 0.20, band 9 fixed 2.60× (the symbol's native size). Size
# therefore encodes speed alongside the colour, exactly as the catalogue does.
# (lo, hi, colour, fixed_scale | None → scale = speed × _SCALE_PER_KN)
_SCALE_FLOOR = 0.40
_SCALE_CEIL = 2.60
_SCALE_PER_KN = 0.20
_SPEED_BANDS = [
    (0.0, 0.5, "#7652E2", _SCALE_FLOOR), (0.5, 1.0, "#4898D3", _SCALE_FLOOR),
    (1.0, 2.0, "#61CBE5", _SCALE_FLOOR), (2.0, 3.0, "#6DBC45", None),
    (3.0, 5.0, "#B4DC00", None), (5.0, 7.0, "#CDC100", None),
    (7.0, 10.0, "#F8A718", None), (10.0, 13.0, "#F7A29D", None),
    (13.0, None, "#FF1E1E", _SCALE_CEIL),
]
# The SCAROW SVGs are 6×11 mm; scale 1.0 = that native size. QGIS SvgMarker size
# is the rendered height, so native height (mm) is the scale-1.0 reference.
_NATIVE_MM = 11.0
_ARROW_SVG = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                          "svg", "repository", "s111_current_arrow.svg")
_NODATA = -9999.0


def _log(msg, level=None):
    try:
        QgsMessageLog.logMessage(msg, "ENCore GIS",
                                 level if level is not None else Qgis.Info)
    except Exception:                                   # noqa: BLE001
        pass


def _gdal():
    from osgeo import gdal
    gdal.UseExceptions()
    return gdal


# ── Detection (folder or .zip), mirroring enc_s102 ───────────────────────────
def is_s111(path):
    """True if GDAL identifies ``path`` as an S-111 dataset (or the 111… HDF5
    name convention)."""
    try:
        drv = _gdal().IdentifyDriver(path)
        if drv is not None and drv.ShortName == "S111":
            return True
    except Exception:                                   # noqa: BLE001
        pass
    b = os.path.basename(path).lower()
    return b.startswith("111") and b.endswith((".h5", ".hdf5"))


def find_s111(folder):
    out = []
    if not folder or not os.path.isdir(folder):
        return out
    for dirpath, _dirs, files in os.walk(folder):
        for fn in files:
            if fn.lower().endswith((".h5", ".hdf5")):
                p = os.path.join(dirpath, fn)
                if is_s111(p):
                    out.append(p)
    return sorted(out)


_ZIP_CACHE = {}


def _extract_s111_from_zip(zip_path):
    if zip_path in _ZIP_CACHE:
        return _ZIP_CACHE[zip_path]
    out = []
    try:
        with zipfile.ZipFile(zip_path) as z:
            members = [m for m in z.namelist()
                       if m.lower().endswith((".h5", ".hdf5"))]
            if members:
                dest = tempfile.mkdtemp(prefix="enc_s111_zip_")
                root = os.path.realpath(dest)
                for m in members:
                    # Preserve the archive's folders: S-111 depth layers share a
                    # basename and are distinguished only by folder (e.g. 0-5/,
                    # 0-10/) — flattening to basename would collide them into one.
                    rel = os.path.normpath(m.replace("\\", "/")).lstrip("/\\")
                    tgt = os.path.join(dest, rel)
                    # Zip-slip guard: reject any member that resolves OUTSIDE the
                    # temp dir (../ traversal OR a drive/UNC-absolute member).
                    if os.path.commonpath(
                            [root, os.path.realpath(tgt)]) != root:
                        _log(f"S-111: skipped unsafe zip member {m!r}",
                             Qgis.Warning)
                        continue
                    os.makedirs(os.path.dirname(tgt) or dest, exist_ok=True)
                    with z.open(m) as src, open(tgt, "wb") as dst:
                        shutil.copyfileobj(src, dst)
                    if is_s111(tgt):
                        out.append(tgt)
    except Exception as e:                              # noqa: BLE001
        _log(f"S-111: reading zip failed ({os.path.basename(zip_path)}): {e}",
             Qgis.Warning)
    out = sorted(out)
    _ZIP_CACHE[zip_path] = out
    return out


def datasets_from_source(source):
    """S-111 datasets for a chart source — a folder (walked) or a .zip (HDF5
    members extracted)."""
    if not source:
        return []
    if os.path.isdir(source):
        return find_s111(source)
    if zipfile.is_zipfile(source):
        return _extract_s111_from_zip(source)
    return []


# ── Time axis ────────────────────────────────────────────────────────────────
def _groups(ds):
    """Sorted [(index0, subdataset_name)] of the Group_NNN time records."""
    subs = ds.GetSubDatasets() if ds is not None else []
    groups = []
    for name, _desc in subs:
        m = re.search(r"Group_(\d+)\s*$", name)
        if m:
            groups.append((int(m.group(1)), name))
    groups.sort()
    return [(i, name) for i, (_n, name) in enumerate(groups)]


def _parse_dt(s):
    for fmt in ("%Y%m%dT%H%M%SZ", "%Y%m%dT%H%M%S", "%Y-%m-%dT%H:%M:%SZ"):
        try:
            return datetime.strptime(str(s).strip(), fmt)
        except (ValueError, TypeError):
            continue
    return None


def time_axis(paths):
    """The UNION time axis over ``paths`` as a sorted list of datetimes. Uses
    dateTimeOfFirstRecord + timeRecordInterval + numberOfTimes (falling back to
    the count of Group_NNN records). Empty list if none readable."""
    gdal = _gdal()
    times = set()
    for p in _as_list(paths):
        try:
            ds = gdal.Open(p)
            md = ds.GetMetadata()
            first = _parse_dt(md.get("dateTimeOfFirstRecord"))
            interval = int(float(md.get("timeRecordInterval") or 0))
            n = int(float(md.get("numberOfTimes") or 0)) or len(_groups(ds))
            if first and interval and n:
                for i in range(n):
                    times.add(first + timedelta(seconds=interval * i))
            elif first:
                times.add(first)
        except Exception as e:                          # noqa: BLE001
            _log(f"S-111 time axis failed ({os.path.basename(p)}): {e}",
                 Qgis.Warning)
    return sorted(times)


def _as_list(paths):
    if isinstance(paths, (list, tuple)):
        return list(paths)
    return [paths]


# ── Current-arrow field for one time ─────────────────────────────────────────
def current_field(paths, when, target_crs, name="S-111 surface currents",
                  step=3):
    """Build a point ``QgsVectorLayer`` of current arrows for the time closest to
    ``when`` (a datetime), portrayed per the S-111 speed-band scale + rotation.
    ``step`` decimates the grid (every Nth node) to keep the field readable.
    Returns the layer, or None if nothing rendered."""
    if Qgis is None:
        return None
    gdal = _gdal()
    src = QgsCoordinateReferenceSystem("EPSG:4326")
    dst = QgsCoordinateReferenceSystem(target_crs)
    if not dst.isValid():
        dst = src
    xform = (None if dst == src
             else QgsCoordinateTransform(src, dst, QgsProject.instance()))
    layer = QgsVectorLayer(f"Point?crs={dst.authid() or 'EPSG:4326'}",
                           name, "memory")
    dp = layer.dataProvider()
    dp.addAttributes([QgsField("speed", QVariant.Double),
                      QgsField("direction", QVariant.Double)])
    layer.updateFields()

    feats = []
    for p in _as_list(paths):
        try:
            ds = gdal.Open(p)
            groups = _groups(ds)
            if not groups:
                continue
            idx = _nearest_index(p, when, len(groups))
            sub = gdal.Open(groups[idx][1])
            gt = sub.GetGeoTransform()
            speed = sub.GetRasterBand(1).ReadAsArray()
            direction = sub.GetRasterBand(2).ReadAsArray()
        except Exception as e:                          # noqa: BLE001
            _log(f"S-111 read failed ({os.path.basename(p)}): {e}", Qgis.Warning)
            continue
        if speed is None or direction is None:
            continue
        ny, nx = speed.shape
        st = max(1, int(step))
        for row in range(0, ny, st):
            for col in range(0, nx, st):
                sp = float(speed[row, col])
                di = float(direction[row, col])
                # Per the PC (main.xsl note 4): no symbol for zero or noData speed.
                if sp <= _NODATA + 1 or di <= _NODATA + 1 or sp <= 0:
                    continue
                x = gt[0] + (col + 0.5) * gt[1] + (row + 0.5) * gt[2]
                y = gt[3] + (col + 0.5) * gt[4] + (row + 0.5) * gt[5]
                g = QgsGeometry.fromPointXY(QgsPointXY(x, y))
                if xform is not None:
                    g.transform(xform)
                ft = QgsFeature(layer.fields())
                ft.setGeometry(g)
                ft.setAttributes([sp, di])
                feats.append(ft)
    if not feats:
        return None
    dp.addFeatures(feats)
    layer.updateExtents()
    layer.setRenderer(_arrow_renderer())
    return layer


def _nearest_index(path, when, n):
    """Index of the Group_NNN closest to datetime ``when`` for this file."""
    if when is None:
        return 0
    try:
        md = _gdal().Open(path).GetMetadata()
        first = _parse_dt(md.get("dateTimeOfFirstRecord"))
        interval = int(float(md.get("timeRecordInterval") or 0))
        if first and interval:
            i = round((when - first).total_seconds() / interval)
            return max(0, min(n - 1, int(i)))
    except Exception:                                   # noqa: BLE001
        pass
    return 0


def _arrow_renderer():
    """The official S-111 9-band arrow renderer: one current-arrow per speed band,
    filled with the band's STEP colour, ROTATED (data-defined) to
    surfaceCurrentDirection, and SIZED per the catalogue's scale rule — fixed
    0.40× (bands 1-3) / 2.60× (band 9), or speed × 0.20× (bands 4-8) of the
    symbol's native size. See _SPEED_BANDS."""
    root = QgsRuleBasedRenderer.Rule(None)
    for lo, hi, colour, fixed in _SPEED_BANDS:
        sl = QgsSvgMarkerSymbolLayer(_ARROW_SVG)
        sl.setFillColor(_qcolor(colour))
        sl.setStrokeColor(_qcolor("#000000"))
        sl.setStrokeWidth(0.2)
        if fixed is not None:
            sl.setSize(fixed * _NATIVE_MM)
        else:
            # scale = surfaceCurrentSpeed × 0.20 (of native) — size encodes speed
            sl.setSize(_SCALE_PER_KN * _NATIVE_MM)   # fallback constant
            sl.setDataDefinedProperty(
                QgsSymbolLayer.PropertySize,
                QgsProperty.fromExpression(
                    f'"speed" * {_SCALE_PER_KN} * {_NATIVE_MM}'))
        sym = QgsMarkerSymbol([sl])
        sym.setDataDefinedAngle(QgsProperty.fromField("direction"))
        expr = (f'"speed" >= {lo}' if hi is None
                else f'"speed" >= {lo} AND "speed" < {hi}')
        root.appendChild(QgsRuleBasedRenderer.Rule(
            sym, filterExp=expr, label=f"{lo}–{hi if hi else '+'} kn"))
    return QgsRuleBasedRenderer(root)


def _qcolor(hex_str):
    from qgis.PyQt.QtGui import QColor
    return QColor(hex_str)
