"""Pipeline instrumentation for debugging the generation flow."""
from __future__ import annotations

import json
import os
import shutil
import time
from dataclasses import asdict, dataclass


@dataclass
class PipelineContext:
    """Accumulates metadata at each pipeline boundary for debugging.

    All fields are Optional because the context is populated incrementally:
    canvas_exporter -> generation_service -> raster_writer.
    """

    # Canvas Export (populated by canvas_exporter.py)
    extent: dict | None = None
    crs_wkt: str | None = None
    crs_authid: str | None = None
    # Rendered extent reprojected to WGS84 (EPSG:4326) client-side as
    # {west, south, east, north}, so the backend stores an exact footprint
    # without needing its own proj library.
    bbox_wgs84: dict | None = None
    # Centroid of the rendered zone in WGS84, computed client-side via QGIS
    # so the server doesn't need its own reprojection library.
    centroid_lat: float | None = None
    centroid_lon: float | None = None
    # Meters per pixel at the centroid, sent to the backend as ground
    # resolution context.
    ground_resolution_m: float | None = None
    export_width: int | None = None
    export_height: int | None = None
    # Sanitized basemap identity (provider + host only, never a file path or
    # auth token) under the rendered zone. Pure context for the internal team
    # map; safe to send.
    basemap: str | None = None
    aspect_ratio: str | None = None
    image_size_bytes: int | None = None
    # Wire format the input was encoded as ('webp' | 'jpeg' | 'png'). Sent to
    # the upload endpoint so the object is signed with a matching content-type.
    input_format: str | None = None

    # Iteration chain: set by plugin.py when the user iterates on a previous
    # result, so the server can silently attach the original input as a
    # reference image to anchor coherence across edits.
    parent_request_id: str | None = None

    # Template attribution: set when the submitted prompt exactly matches a
    # curated preset (after whitespace normalization). Lets the dashboard
    # distinguish "ran Top Pick X" from "wrote custom prompt".
    template_id: str | None = None
    template_name: str | None = None
    # Vectorize hints carried from the template metadata (catalog), so the
    # post-generation result panel can suggest extracting class regions
    # with the color already pre-filled. vector_color is the single-class
    # hex string; vector_classes is the multi-class list (each entry has
    # color + label). Either may be None.
    vector_color: str | None = None
    vector_classes: list[dict] | None = None

    # Submit (populated by generation_service.py)
    request_id: str | None = None
    submitted_resolution: str | None = None
    submitted_aspect_ratio: str | None = None
    submit_timestamp: float | None = None
    credit_cost: int | None = None
    estimated_time_seconds: float | None = None
    max_wait_seconds: float | None = None

    # Poll (populated by generation_service.py)
    poll_count: int | None = None
    total_wait_seconds: float | None = None
    final_status: str | None = None

    # Download (populated by generation_service.py)
    received_image_width: int | None = None
    received_image_height: int | None = None
    received_size_bytes: int | None = None

    # Write (populated by raster_writer.py)
    output_path: str | None = None
    geotransform: tuple | None = None
    output_bands: int | None = None
    output_dimensions: tuple[int, int] | None = None
    crop_offsets: tuple[int, int, int, int] | None = None  # x, y, w, h

    # Refund safety: set to True once the worker has asked the server to
    # refund this generation, so download_failed + write_error paths can't
    # double-refund.
    refund_emitted: bool = False

    def validate(self) -> list[str]:
        """Check boundary consistency. Returns list of warning strings."""
        warnings = []
        if (
            self.aspect_ratio
            and self.submitted_aspect_ratio  # noqa: W503
            and self.aspect_ratio != self.submitted_aspect_ratio  # noqa: W503
        ):
            warnings.append(
                f"Aspect ratio mismatch: export={self.aspect_ratio}, "
                f"submitted={self.submitted_aspect_ratio}"
            )
        if (
            self.export_width
            and self.received_image_width  # noqa: W503
            and self.export_width != self.received_image_width  # noqa: W503
        ):
            warnings.append(
                f"Width mismatch: sent={self.export_width}, "
                f"received={self.received_image_width}"
            )
        if (
            self.export_height
            and self.received_image_height  # noqa: W503
            and self.export_height != self.received_image_height  # noqa: W503
        ):
            warnings.append(
                f"Height mismatch: sent={self.export_height}, "
                f"received={self.received_image_height}"
            )
        return warnings

    def safe_log_summary(self) -> str:
        """Production-safe log summary. No URLs, API keys, or model names."""
        parts = []
        if self.export_width and self.export_height:
            parts.append(f"Export: {self.export_width}x{self.export_height}px")
        if self.aspect_ratio:
            parts.append(f"ratio={self.aspect_ratio}")
        if self.request_id:
            parts.append(f"request_id={self.request_id}")
        if self.submitted_resolution:
            parts.append(f"resolution={self.submitted_resolution}")
        if self.credit_cost is not None:
            parts.append(f"credits={self.credit_cost}")
        if self.max_wait_seconds:
            parts.append(f"max_wait={int(self.max_wait_seconds)}s")
        if self.received_image_width and self.received_image_height:
            parts.append(
                f"Result: {self.received_image_width}x{self.received_image_height}px"
            )
        if self.received_size_bytes:
            parts.append(f"{self.received_size_bytes // 1024}KB")
        if self.output_path:
            parts.append(f"Output: {self.output_path}")
        return " | ".join(parts)


def save_debug_artifacts(
    ctx: PipelineContext,
    sent_png: bytes | None,
    received_png: bytes | None,
    plugin_dir: str,
    max_runs: int = 20,
    context_images: list[bytes] | None = None,
    guidance_png: bytes | None = None,
    guidance_format: str | None = None,
) -> str | None:
    """Save debug artifacts to .debug/{timestamp}/. Returns path or None."""
    debug_dir = os.path.join(plugin_dir, ".debug")
    run_dir = os.path.join(debug_dir, str(int(time.time())))
    os.makedirs(run_dir, exist_ok=True)

    if sent_png:
        # The sent input is now a compressed format (webp/jpeg), not PNG; name
        # the artifact by its real format so it opens correctly.
        sent_ext = {"webp": "webp", "jpeg": "jpg"}.get(ctx.input_format or "", "png")
        with open(os.path.join(run_dir, f"sent.{sent_ext}"), "wb") as f:
            f.write(sent_png)
    if received_png:
        with open(os.path.join(run_dir, "received.png"), "wb") as f:
            f.write(received_png)

    # The markup-overlay guidance image, when the user drew annotations. This
    # is the original zone with the marks on top; the main image (sent.*) is
    # clean. Saving both is the quickest way to confirm the split is correct.
    if guidance_png:
        guidance_ext = {"webp": "webp", "jpeg": "jpg"}.get(guidance_format or "", "png")
        with open(os.path.join(run_dir, f"guidance.{guidance_ext}"), "wb") as f:
            f.write(guidance_png)

    if context_images:
        for idx, data in enumerate(context_images, start=1):
            with open(os.path.join(run_dir, f"context_{idx}.jpg"), "wb") as f:
                f.write(data)

    # Save both sent and received as GeoTIFFs for visual alignment check.
    # Load sent.tif + received.tif in QGIS to see which one is offset.
    if ctx.extent and ctx.crs_wkt:
        _save_debug_geotiff(
            run_dir, "sent.tif", sent_png, ctx.extent, ctx.crs_wkt
        )
        _save_debug_geotiff(
            run_dir, "received.tif", received_png, ctx.extent, ctx.crs_wkt
        )

    ctx_dict = {}
    for k, v in asdict(ctx).items():
        if v is not None:
            ctx_dict[k] = v
    with open(os.path.join(run_dir, "context.json"), "w", encoding="utf-8") as f:
        json.dump(ctx_dict, f, indent=2, default=str)

    _cleanup_old_runs(debug_dir, max_runs)
    return run_dir


def _save_debug_geotiff(
    run_dir: str,
    filename: str,
    image_bytes: bytes | None,
    extent: dict,
    crs_wkt: str,
):
    """Write a debug GeoTIFF from raw PNG bytes + extent. Fails silently."""
    if not image_bytes:
        return
    try:
        from osgeo import gdal, osr

        temp_png = os.path.join(run_dir, f"_tmp_{filename}.png")
        tif_path = os.path.join(run_dir, filename)
        with open(temp_png, "wb") as f:
            f.write(image_bytes)
        src = gdal.Open(temp_png)
        if src is None:
            return
        w, h, bands = src.RasterXSize, src.RasterYSize, min(src.RasterCount, 3)
        ext_w = extent["xmax"] - extent["xmin"]
        ext_h = extent["ymax"] - extent["ymin"]
        drv = gdal.GetDriverByName("GTiff")
        dst = drv.Create(tif_path, w, h, bands, gdal.GDT_Byte)
        dst.SetGeoTransform((
            extent["xmin"], ext_w / w, 0,
            extent["ymax"], 0, -(ext_h / h),
        ))
        srs = osr.SpatialReference()
        srs.ImportFromWkt(crs_wkt)
        dst.SetProjection(srs.ExportToWkt())
        for i in range(1, bands + 1):
            dst.GetRasterBand(i).WriteArray(src.GetRasterBand(i).ReadAsArray())
        dst.FlushCache()
        dst = None
        src = None
        os.remove(temp_png)
    except Exception:
        pass  # nosec B110


def _cleanup_old_runs(debug_dir: str, max_runs: int):
    """Delete oldest debug runs beyond max_runs."""
    if not os.path.isdir(debug_dir):
        return
    runs = sorted(
        [d for d in os.listdir(debug_dir) if os.path.isdir(os.path.join(debug_dir, d))],
        reverse=True,
    )
    for old_run in runs[max_runs:]:
        shutil.rmtree(os.path.join(debug_dir, old_run), ignore_errors=True)
