#!/usr/bin/env python3
"""Build a single XYZ PNG tile layer from KML/KMZ GroundOverlay inputs."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import shutil
import subprocess
import sys
import tempfile
import uuid
import zipfile
from datetime import UTC, datetime
from pathlib import Path
from urllib.parse import unquote, urlparse
from xml.etree import ElementTree as ET

from PIL import Image


PIPELINE_VERSION = "1.0.0"
WORLD_WIDTH_METERS = 40075016.68557849
MAX_CONFIG_ZOOM = 22


class SourceError(RuntimeError):
    """An input-specific error which must not abort the complete build."""


def run(command: list[str]) -> None:
    completed = subprocess.run(command, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if completed.returncode:
        detail = completed.stderr.strip() or completed.stdout.strip() or "comando sem detalhe"
        raise SourceError(f"{' '.join(command[:2])} falhou: {detail}")


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def local_name(element: ET.Element) -> str:
    return element.tag.rsplit("}", 1)[-1]


def first_child(parent: ET.Element, name: str) -> ET.Element | None:
    for child in parent.iter():
        if local_name(child) == name:
            return child
    return None


def child_text(parent: ET.Element, name: str) -> str | None:
    child = first_child(parent, name)
    return child.text.strip() if child is not None and child.text else None


def mercator(lon: float, lat: float) -> tuple[float, float]:
    if not -85.05112878 < lat < 85.05112878:
        raise SourceError(f"latitude fora do limite do Web Mercator: {lat}")
    radius = 6378137.0
    return (
        radius * math.radians(lon),
        radius * math.log(math.tan(math.pi / 4.0 + math.radians(lat) / 2.0)),
    )


def inverse_mercator(x: float, y: float) -> tuple[float, float]:
    radius = 6378137.0
    return math.degrees(x / radius), math.degrees(2.0 * math.atan(math.exp(y / radius)) - math.pi / 2.0)


def rotated_corners(box: dict[str, float]) -> list[tuple[float, float]]:
    """Return UL, UR, LR, LL corners in EPSG:3857 after KML rotation."""
    raw = [
        mercator(box["west"], box["north"]),
        mercator(box["east"], box["north"]),
        mercator(box["east"], box["south"]),
        mercator(box["west"], box["south"]),
    ]
    center_x = sum(point[0] for point in raw) / 4.0
    center_y = sum(point[1] for point in raw) / 4.0
    radians = math.radians(box["rotation"])
    cosine, sine = math.cos(radians), math.sin(radians)
    return [
        (
            center_x + (x - center_x) * cosine - (y - center_y) * sine,
            center_y + (x - center_x) * sine + (y - center_y) * cosine,
        )
        for x, y in raw
    ]


def polygon_area(points: list[tuple[float, float]]) -> float:
    return abs(
        sum(points[index][0] * points[(index + 1) % len(points)][1] - points[(index + 1) % len(points)][0] * points[index][1]
            for index in range(len(points)))
    ) / 2.0


def bounds(points: list[tuple[float, float]]) -> list[float]:
    return [min(point[0] for point in points), min(point[1] for point in points), max(point[0] for point in points), max(point[1] for point in points)]


def validate_zoom(value: object, label: str) -> int | None:
    if value is None:
        return None
    if not isinstance(value, int) or isinstance(value, bool) or not 0 <= value <= MAX_CONFIG_ZOOM:
        raise ValueError(f"{label} deve ser inteiro entre 0 e {MAX_CONFIG_ZOOM} ou null")
    return value


def read_config(path: Path) -> dict:
    if not path.exists():
        return {"global": {"minzoom": None, "maxzoom": None}, "sources": {}}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise RuntimeError(f"configuração inválida em {path}: {error}") from error
    if not isinstance(payload, dict) or not isinstance(payload.get("global", {}), dict) or not isinstance(payload.get("sources", {}), dict):
        raise RuntimeError("maps.json deve conter objetos 'global' e 'sources'")
    global_config = payload.get("global", {})
    result = {
        "global": {
            "minzoom": validate_zoom(global_config.get("minzoom"), "global.minzoom"),
            "maxzoom": validate_zoom(global_config.get("maxzoom"), "global.maxzoom"),
        },
        "sources": {},
    }
    for name, item in payload.get("sources", {}).items():
        if not isinstance(name, str) or not isinstance(item, dict):
            raise RuntimeError("cada entrada em sources deve ter caminho textual e objeto de zoom")
        result["sources"][name.replace("\\", "/")] = {
            "minzoom": validate_zoom(item.get("minzoom"), f"sources.{name}.minzoom"),
            "maxzoom": validate_zoom(item.get("maxzoom"), f"sources.{name}.maxzoom"),
        }
    if result["global"]["minzoom"] is not None and result["global"]["maxzoom"] is not None and result["global"]["minzoom"] > result["global"]["maxzoom"]:
        raise RuntimeError("global.minzoom não pode exceder global.maxzoom")
    return result


def effective_zooms(auto_max: int, relative_source: str, config: dict) -> tuple[int, int, bool]:
    auto_min = max(6, min(10, auto_max - 7))
    override = config["sources"].get(relative_source, {})
    minimum = override.get("minzoom")
    maximum = override.get("maxzoom")
    if minimum is None:
        minimum = config["global"]["minzoom"]
    if maximum is None:
        maximum = config["global"]["maxzoom"]
    minimum = auto_min if minimum is None else minimum
    maximum = auto_max if maximum is None else maximum
    if minimum > maximum:
        raise SourceError(f"minzoom ({minimum}) excede maxzoom ({maximum})")
    return minimum, maximum, maximum > auto_max


def safe_extract_kmz(archive: Path, destination: Path) -> Path:
    try:
        with zipfile.ZipFile(archive) as zipped:
            for member in zipped.infolist():
                target = (destination / member.filename).resolve()
                if target != destination.resolve() and destination.resolve() not in target.parents:
                    raise SourceError(f"KMZ contém caminho inseguro: {member.filename}")
                if member.is_dir():
                    target.mkdir(parents=True, exist_ok=True)
                else:
                    target.parent.mkdir(parents=True, exist_ok=True)
                    with zipped.open(member) as source, target.open("wb") as output:
                        shutil.copyfileobj(source, output)
    except zipfile.BadZipFile as error:
        raise SourceError(f"KMZ inválido: {error}") from error
    candidates = sorted(destination.rglob("*.kml"), key=lambda item: item.as_posix().lower())
    preferred = [item for item in candidates if item.relative_to(destination).as_posix().lower() == "doc.kml"]
    root_level = [item for item in candidates if item.parent == destination]
    if preferred:
        return preferred[0]
    if root_level:
        return sorted(root_level, key=lambda item: item.name.lower())[0]
    if candidates:
        return candidates[0]
    raise SourceError("KMZ não contém KML")


def resolve_image(href: str, kml_directory: Path, allowed_root: Path) -> Path:
    parsed = urlparse(href)
    if parsed.scheme or parsed.netloc or href.startswith("/"):
        raise SourceError("Icon/href externo ou absoluto não é suportado")
    candidate = (kml_directory / unquote(parsed.path)).resolve()
    root = allowed_root.resolve()
    if candidate != root and root not in candidate.parents:
        raise SourceError("Icon/href sai da árvore permitida")
    if not candidate.is_file():
        raise SourceError(f"imagem referenciada não encontrada: {href}")
    return candidate


def parse_overlays(kml_file: Path, allowed_root: Path) -> list[dict]:
    try:
        document = ET.parse(kml_file)
    except (OSError, ET.ParseError) as error:
        raise SourceError(f"KML inválido: {error}") from error
    overlays: list[dict] = []
    for index, element in enumerate(item for item in document.getroot().iter() if local_name(item) == "GroundOverlay"):
        icon = first_child(element, "Icon")
        lat_lon_box = first_child(element, "LatLonBox")
        if icon is None or lat_lon_box is None:
            raise SourceError(f"GroundOverlay #{index + 1} sem Icon ou LatLonBox")
        href = child_text(icon, "href")
        if not href:
            raise SourceError(f"GroundOverlay #{index + 1} sem Icon/href")
        try:
            box = {name: float(child_text(lat_lon_box, name) or "") for name in ("north", "south", "east", "west")}
            box["rotation"] = float(child_text(lat_lon_box, "rotation") or "0")
        except ValueError as error:
            raise SourceError(f"GroundOverlay #{index + 1} possui LatLonBox inválido") from error
        if not box["north"] > box["south"] or not box["east"] > box["west"]:
            raise SourceError(f"GroundOverlay #{index + 1} possui LatLonBox sem extensão positiva")
        overlays.append({
            "index": index + 1,
            "name": child_text(element, "name") or f"overlay-{index + 1}",
            "href": href,
            "image": str(resolve_image(href, kml_file.parent, allowed_root)),
            "latlonbox": box,
        })
    if not overlays:
        raise SourceError("KML não contém GroundOverlay")
    return overlays


def gdal_size(path: Path) -> tuple[int, int]:
    completed = subprocess.run(["gdalinfo", "-json", str(path)], text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if completed.returncode:
        raise SourceError(completed.stderr.strip() or "gdalinfo não conseguiu ler a imagem")
    try:
        size = json.loads(completed.stdout)["size"]
        return int(size[0]), int(size[1])
    except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error:
        raise SourceError("gdalinfo não retornou dimensões da imagem") from error


def warp_overlay(overlay: dict, output: Path, workspace: Path) -> dict:
    width, height = gdal_size(Path(overlay["image"]))
    corners = rotated_corners(overlay["latlonbox"])
    vrt = workspace / f"overlay-{overlay['index']}.vrt"
    gcp_args: list[str] = []
    for pixel_x, pixel_y, point in zip((0, width, width, 0), (0, 0, height, height), corners):
        gcp_args.extend(["-gcp", str(pixel_x), str(pixel_y), f"{point[0]:.12f}", f"{point[1]:.12f}"])
    run(["gdal_translate", "-q", "-of", "VRT", "-a_srs", "EPSG:3857", *gcp_args, overlay["image"], str(vrt)])
    run([
        "gdalwarp", "-q", "-overwrite", "-of", "GTiff", "-t_srs", "EPSG:3857", "-order", "1", "-dstalpha",
        "-r", "bilinear", "-multi", "-co", "TILED=YES", "-co", "COMPRESS=DEFLATE", "-co", "BIGTIFF=IF_SAFER",
        str(vrt), str(output),
    ])
    completed = subprocess.run(["gdalinfo", "-json", str(output)], text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if completed.returncode:
        raise SourceError(completed.stderr.strip() or "gdalinfo não conseguiu ler o raster reprojetado")
    info = json.loads(completed.stdout)
    transform = info.get("geoTransform")
    if not transform:
        raise SourceError("raster reprojetado não possui geotransformação")
    resolution = max(abs(float(transform[1])), abs(float(transform[5])))
    auto_max = max(0, min(MAX_CONFIG_ZOOM, math.floor(math.log2(WORLD_WIDTH_METERS / (256.0 * resolution)))))
    wgs_corners = [inverse_mercator(*point) for point in corners]
    overlay.update({
        "warped": str(output),
        "area_m2": polygon_area(corners),
        "bbox_3857": bounds(corners),
        "bbox_wgs84": bounds(wgs_corners),
        "pixel_resolution_m": resolution,
        "auto_maxzoom": auto_max,
    })
    return overlay


def source_cache_dir(cache_root: Path, relative_source: str) -> Path:
    return cache_root / "sources" / hashlib.sha256(relative_source.encode("utf-8")).hexdigest()[:24]


def read_cached_source(cache_dir: Path, source_hash: str) -> dict | None:
    descriptor_path = cache_dir / "descriptor.json"
    if not descriptor_path.exists():
        return None
    try:
        descriptor = json.loads(descriptor_path.read_text(encoding="utf-8"))
        if descriptor.get("pipeline_version") != PIPELINE_VERSION or descriptor.get("source_sha256") != source_hash:
            return None
        if not all(Path(item["warped"]).is_file() for item in descriptor.get("overlays", [])):
            return None
        return descriptor
    except (OSError, ValueError, KeyError):
        return None


def process_source(source: Path, relative_source: str, cache_root: Path, source_hash: str) -> dict:
    cache_dir = source_cache_dir(cache_root, relative_source)
    cached = read_cached_source(cache_dir, source_hash)
    if cached is not None:
        cached["cache_hit"] = True
        return cached
    if cache_dir.exists():
        shutil.rmtree(cache_dir)
    cache_dir.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="arma-tiles-", dir=cache_root) as temp_dir:
        temporary = Path(temp_dir)
        if source.suffix.lower() == ".kmz":
            extracted_root = temporary / "kmz"
            extracted_root.mkdir()
            kml_file = safe_extract_kmz(source, extracted_root)
            allowed_root = extracted_root
        else:
            kml_file = source
            allowed_root = source.parent
        overlays = parse_overlays(kml_file, allowed_root)
        processed: list[dict] = []
        for overlay in overlays:
            warped = cache_dir / f"overlay-{overlay['index']}.tif"
            processed_overlay = warp_overlay(overlay, warped, temporary)
            processed_overlay["source"] = relative_source
            processed.append(processed_overlay)
    descriptor = {
        "pipeline_version": PIPELINE_VERSION,
        "source": relative_source,
        "source_sha256": source_hash,
        "overlays": processed,
        "cache_hit": False,
    }
    (cache_dir / "descriptor.json").write_text(json.dumps(descriptor, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return descriptor


def alpha_is_empty(image: Image.Image) -> bool:
    return image.getchannel("A").getbbox() is None


def compose_layer(source_tiles: Path, final_tiles: Path) -> None:
    for tile in sorted(source_tiles.rglob("*.png")):
        relative = tile.relative_to(source_tiles)
        destination = final_tiles / relative
        with Image.open(tile) as opened:
            layer = opened.convert("RGBA")
            layer.load()
        if alpha_is_empty(layer):
            continue
        destination.parent.mkdir(parents=True, exist_ok=True)
        if destination.exists():
            with Image.open(destination) as opened:
                base = opened.convert("RGBA")
                base.load()
            base.alpha_composite(layer)
            result = base
        else:
            result = layer
        if alpha_is_empty(result):
            destination.unlink(missing_ok=True)
        else:
            result.save(destination, format="PNG", optimize=False)


def tile_overlay(overlay: dict, work_root: Path, final_tiles: Path, processes: int) -> None:
    layer_root = work_root / f"layer-{overlay['priority']:04d}-{overlay['index']}"
    help_result = subprocess.run(["gdal2tiles.py", "--help"], text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    legacy = ["--legacy"] if "--legacy" in help_result.stdout else []
    run([
        "gdal2tiles.py", *legacy, "-q", "--profile=mercator", "--xyz", "--exclude", "--webviewer=none", "--resampling=bilinear",
        f"--zoom={overlay['minzoom']}-{overlay['maxzoom']}", f"--processes={processes}", overlay["warped"], str(layer_root),
    ])
    compose_layer(layer_root, final_tiles)
    shutil.rmtree(layer_root, ignore_errors=True)


def publish(stage: Path, data_dir: Path) -> None:
    current = data_dir / "tiles"
    backup = data_dir / ".tiles-previous"
    if backup.exists():
        shutil.rmtree(backup)
    if current.exists():
        current.rename(backup)
    stage.rename(current)
    if backup.exists():
        shutil.rmtree(backup)


def source_record(descriptor: dict, config: dict) -> dict:
    overlays = []
    for item in descriptor["overlays"]:
        minimum, maximum, expands = effective_zooms(item["auto_maxzoom"], descriptor["source"], config)
        item["minzoom"] = minimum
        item["maxzoom"] = maximum
        item["configured_above_native"] = expands
        overlays.append({
            "index": item["index"], "name": item["name"], "href": item["href"], "latlonbox": item["latlonbox"],
            "bbox_wgs84": item["bbox_wgs84"], "bbox_3857": item["bbox_3857"], "area_m2": item["area_m2"],
            "pixel_resolution_m": item["pixel_resolution_m"], "auto_maxzoom": item["auto_maxzoom"],
            "minzoom": minimum, "maxzoom": maximum, "configured_above_native": expands,
        })
    all_points = [point for item in descriptor["overlays"] for point in [
        (item["bbox_wgs84"][0], item["bbox_wgs84"][1]), (item["bbox_wgs84"][2], item["bbox_wgs84"][3])
    ]]
    all_projected_points = [point for item in descriptor["overlays"] for point in [
        (item["bbox_3857"][0], item["bbox_3857"][1]), (item["bbox_3857"][2], item["bbox_3857"][3])
    ]]
    return {
        "source": descriptor["source"], "sha256": descriptor["source_sha256"], "cache_hit": descriptor["cache_hit"],
        "ground_overlays_processed": len(overlays), "area_m2_total": sum(item["area_m2"] for item in overlays),
        "bbox_wgs84": bounds(all_points), "bbox_3857": bounds(all_projected_points), "overlays": overlays,
    }


def build(arguments: argparse.Namespace) -> None:
    source_dir = Path(arguments.source_dir).resolve()
    data_dir = Path(arguments.data_dir).resolve()
    config = read_config(Path(arguments.config).resolve())
    if not source_dir.is_dir():
        raise RuntimeError(f"diretório source inexistente: {source_dir}")
    data_dir.mkdir(parents=True, exist_ok=True)
    cache_root = data_dir / "cache"
    cache_root.mkdir(parents=True, exist_ok=True)
    files = sorted((path for path in source_dir.rglob("*") if path.is_file() and path.suffix.lower() in {".kml", ".kmz"}), key=lambda item: item.as_posix().lower())
    selected = set(arguments.source_filter or [])
    errors: list[dict] = []
    records: list[dict] = []
    descriptors: list[dict] = []
    for source in files:
        relative = source.relative_to(source_dir).as_posix()
        if selected and relative not in selected:
            continue
        try:
            descriptor = process_source(source, relative, cache_root, sha256_file(source))
            record = source_record(descriptor, config)
            records.append(record)
            descriptors.append(descriptor)
        except (SourceError, OSError, ValueError) as error:
            errors.append({"source": relative, "stage": "processamento", "reason": str(error)})
    overlays = [item for descriptor in descriptors for item in descriptor["overlays"]]
    overlays.sort(key=lambda item: (item["area_m2"], item["source"], item["index"]))
    for priority, overlay in enumerate(overlays, start=1):
        overlay["priority"] = priority
    for record in records:
        by_index = {item["index"]: item for item in record["overlays"]}
        for overlay in (item for item in overlays if item["source"] == record["source"]):
            by_index[overlay["index"]]["priority"] = overlay["priority"]
        record["priority_effective"] = min((item["priority"] for item in record["overlays"]), default=None)

    stage = data_dir / f".tiles-staging-{uuid.uuid4().hex}"
    work_root = data_dir / f".tiles-work-{uuid.uuid4().hex}"
    stage.mkdir(parents=True)
    work_root.mkdir(parents=True)
    try:
        # Larger areas are drawn first. For equal areas, reverse lexical order is drawn first so the documented
        # alphabetical winner is painted last.
        for overlay in sorted(overlays, key=lambda item: (item["area_m2"], item["source"], item["index"]), reverse=True):
            try:
                tile_overlay(overlay, work_root, stage, arguments.processes)
            except SourceError as error:
                errors.append({"source": overlay["source"], "overlay": overlay["index"], "stage": "tile", "reason": str(error)})
        global_min = min((item["minzoom"] for item in overlays), default=None)
        global_max = max((item["maxzoom"] for item in overlays), default=None)
        metadata = {
            "pipeline_version": PIPELINE_VERSION,
            "generated_at": datetime.now(UTC).isoformat(),
            "minzoom": global_min,
            "maxzoom": global_max,
            "processed_sources": records,
            "sources_with_error": errors,
            "source_count": len(records),
            "error_count": len(errors),
        }
        (stage / "metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        publish(stage, data_dir)
    except Exception:
        shutil.rmtree(stage, ignore_errors=True)
        raise
    finally:
        shutil.rmtree(work_root, ignore_errors=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-dir", default=os.environ.get("SOURCE_DIR", "source"))
    parser.add_argument("--data-dir", default=os.environ.get("DATA_DIR", "."))
    parser.add_argument("--config", default=os.environ.get("MAPS_CONFIG", "config/maps.json"))
    parser.add_argument("--source-filter", action="append", help="caminho relativo exato sob source/; repetível")
    parser.add_argument("--processes", type=int, default=max(1, int(os.environ.get("GDAL2TILES_PROCESSES", "1"))))
    return parser.parse_args()


if __name__ == "__main__":
    try:
        build(parse_args())
    except Exception as error:
        print(f"erro fatal do pipeline: {error}", file=sys.stderr)
        sys.exit(1)