import asyncio
import json
import logging
import os
import re
import tarfile
import tempfile
from datetime import datetime, timezone
from pathlib import Path

import asyncpg
import h5py
import numpy as np
from PIL import Image

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("radar-renderer")

DATA_DIR = Path(os.getenv("DATA_DIR", "/data"))
RAW_RADAR_DIR = DATA_DIR / "raw" / "radar"
RENDERED_RADAR_DIR = DATA_DIR / "processed" / "radar"
RADAR_PRODUCT = os.getenv("RADAR_PRODUCT", "rv")
RENDER_INTERVAL_SECONDS = int(os.getenv("RADAR_RENDER_INTERVAL_SECONDS", "300"))
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://hexawetter:hexawetter@postgres:5432/hexawetter")
RADAR_FILE_RE = re.compile(r"^[A-Za-z0-9_.-]+\.(?:tar|tar\.bz2|bz2|gz)$", re.IGNORECASE)

# DWD RV composite extent (EPSG:4326 approximate)
DEFAULT_BOUNDS = {
    "south": 47.0,
    "west": 3.0,
    "north": 56.0,
    "east": 16.0,
}

# Precipitation colormap (mm/h)
COLOR_STOPS = [
    (0.0, (0, 0, 0, 0)),
    (0.1, (100, 180, 255, 80)),
    (0.5, (0, 120, 255, 140)),
    (1.0, (0, 200, 80, 170)),
    (2.0, (255, 255, 0, 200)),
    (5.0, (255, 140, 0, 220)),
    (10.0, (255, 0, 0, 230)),
    (20.0, (180, 0, 120, 240)),
    (50.0, (120, 0, 80, 250)),
]


def apply_colormap(values: np.ndarray) -> np.ndarray:
    rgba = np.zeros((*values.shape, 4), dtype=np.uint8)
    mask = np.isfinite(values) & (values > 0)
    if not np.any(mask):
        return rgba
    vmax = max(float(np.nanpercentile(values[mask], 99)), 1.0)
    normalized = np.clip(values / vmax, 0, 1)
    stops = COLOR_STOPS
    for idx in range(len(stops) - 1):
        low_val, low_color = stops[idx]
        high_val, high_color = stops[idx + 1]
        if high_val <= low_val:
            continue
        band = mask & (values >= low_val) & (values < high_val)
        if not np.any(band):
            continue
        ratio = (values[band] - low_val) / (high_val - low_val)
        for channel in range(4):
            rgba[..., channel][band] = (
                low_color[channel] + ratio * (high_color[channel] - low_color[channel])
            ).astype(np.uint8)
    top_band = mask & (values >= stops[-1][0])
    if np.any(top_band):
        for channel in range(4):
            rgba[..., channel][top_band] = stops[-1][1][channel]
    return rgba


def read_hdf5_dataset(path: Path) -> np.ndarray | None:
    try:
        with h5py.File(path, "r") as handle:
            if "dataset1/data1/data" in handle:
                data = handle["dataset1/data1/data"][:]
                what = handle["dataset1/what"].attrs
                nodata = what.get("nodata", -9999)
                data = np.where(data == nodata, np.nan, data)
                return data.astype(np.float32)
            for key in handle.keys():
                if isinstance(handle[key], h5py.Group) and "data" in handle[key]:
                    subgroup = handle[key]
                    for subkey in subgroup.keys():
                        if "data" in subgroup[subkey]:
                            return np.array(subgroup[subkey]["data"], dtype=np.float32)
    except Exception as exc:
        logger.warning("HDF5 read failed for %s: %s", path.name, exc)
    return None


def extract_archive(archive_path: Path, temp_dir: Path) -> list[Path]:
    extracted: list[Path] = []
    if archive_path.suffix == ".bz2" and archive_path.name.endswith(".tar.bz2"):
        mode = "r:bz2"
    elif archive_path.suffix == ".gz":
        mode = "r:gz"
    else:
        mode = "r"
    with tarfile.open(archive_path, mode) as tar:
        tar.extractall(temp_dir, filter="data")
    for path in temp_dir.rglob("*"):
        if path.is_file():
            extracted.append(path)
    return extracted


def render_png(data: np.ndarray, output_path: Path) -> dict:
    rgba = apply_colormap(data)
    image = Image.fromarray(rgba, mode="RGBA")
    image = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path, format="PNG", optimize=True)
    return DEFAULT_BOUNDS


async def upsert_meta(
    pool: asyncpg.Pool,
    filename: str,
    file_size: int,
    modified_at: datetime,
    render_status: str,
    render_path: str | None = None,
    render_bounds: dict | None = None,
    render_error: str | None = None,
) -> None:
    await pool.execute(
        """
        INSERT INTO radar_file_meta (
            product, filename, file_size, modified_at, render_status,
            render_path, render_bounds, render_error, updated_at
        ) VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, NOW())
        ON CONFLICT (filename) DO UPDATE SET
            file_size = EXCLUDED.file_size,
            modified_at = EXCLUDED.modified_at,
            render_status = EXCLUDED.render_status,
            render_path = EXCLUDED.render_path,
            render_bounds = EXCLUDED.render_bounds,
            render_error = EXCLUDED.render_error,
            updated_at = NOW()
        """,
        RADAR_PRODUCT,
        filename,
        file_size,
        modified_at,
        render_status,
        render_path,
        json.dumps(render_bounds) if render_bounds else None,
        render_error,
    )


async def process_file(pool: asyncpg.Pool, archive_path: Path) -> None:
    filename = archive_path.name
    modified_at = datetime.fromtimestamp(archive_path.stat().st_mtime, timezone.utc)
    output_path = RENDERED_RADAR_DIR / RADAR_PRODUCT / f"{filename}.png"
    if output_path.exists() and output_path.stat().st_size > 0:
        await upsert_meta(
            pool,
            filename,
            archive_path.stat().st_size,
            modified_at,
            "done",
            str(output_path),
            DEFAULT_BOUNDS,
        )
        return

    await upsert_meta(pool, filename, archive_path.stat().st_size, modified_at, "processing")
    try:
        with tempfile.TemporaryDirectory() as temp_name:
            temp_dir = Path(temp_name)
            extracted = extract_archive(archive_path, temp_dir)
            data = None
            for candidate in extracted:
                if candidate.suffix.lower() in {".h5", ".hdf5"} or "RV" in candidate.name.upper():
                    data = read_hdf5_dataset(candidate)
                    if data is not None:
                        break
            if data is None:
                for candidate in extracted:
                    data = read_hdf5_dataset(candidate)
                    if data is not None:
                        break
            if data is None:
                raise RuntimeError("No renderable HDF5 dataset found in archive")

            bounds = render_png(data, output_path)
            await upsert_meta(
                pool,
                filename,
                archive_path.stat().st_size,
                modified_at,
                "done",
                str(output_path),
                bounds,
            )
            logger.info("rendered %s -> %s", filename, output_path)
    except Exception as exc:
        await upsert_meta(
            pool,
            filename,
            archive_path.stat().st_size,
            modified_at,
            "error",
            render_error=str(exc),
        )
        logger.error("render failed for %s: %s", filename, exc)


async def render_pending() -> None:
    target_dir = RAW_RADAR_DIR / "composite" / RADAR_PRODUCT
    target_dir.mkdir(parents=True, exist_ok=True)
    RENDERED_RADAR_DIR.mkdir(parents=True, exist_ok=True)

    pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=3)
    try:
        archives = sorted(
            [p for p in target_dir.iterdir() if p.is_file() and RADAR_FILE_RE.fullmatch(p.name)],
            key=lambda p: p.stat().st_mtime,
        )[-12:]
        for archive_path in archives:
            await process_file(pool, archive_path)
    finally:
        await pool.close()


async def main() -> None:
    logger.info("radar renderer started (product=%s)", RADAR_PRODUCT)
    while True:
        try:
            await render_pending()
        except Exception as exc:
            logger.exception("render cycle failed: %s", exc)
        await asyncio.sleep(RENDER_INTERVAL_SECONDS)


if __name__ == "__main__":
    asyncio.run(main())
