#!/usr/bin/env python3
"""MPU client: request 167.png by default and downlink one tiny preview.

Uses only Python's standard library and the telemetry HTTP API. Fetches safe
167.png first, with one alternate listed PNG if fetching it fails. Sources are
never written to disk. Non-interlaced PNGs are supported.
"""
import argparse
import hashlib
import json
import os
import signal
import struct
import sys
import time
import urllib.request
import zlib
from datetime import datetime, timezone
from pathlib import Path

PNG = b"\x89PNG\r\n\x1a\n"
MAX_IMAGE = 16 * 1024 * 1024


def safe_name(name):
    return (
        isinstance(name, str)
        and 0 < len(name) <= 128
        and name not in (".", "..")
        and "/" not in name
        and "\\" not in name
        and "\x00" not in name
        and name.endswith(".png")
    )


def png_chunks(data):
    if not data.startswith(PNG):
        raise ValueError("invalid PNG signature")
    position = 8
    while position + 12 <= len(data):
        length, kind = struct.unpack_from(">I4s", data, position)
        end = position + 12 + length
        if end > len(data):
            raise ValueError("truncated PNG chunk")
        content = data[position + 8 : end - 4]
        crc = struct.unpack_from(">I", data, end - 4)[0]
        if zlib.crc32(kind + content) & 0xFFFFFFFF != crc:
            raise ValueError("PNG chunk CRC mismatch")
        yield kind, content
        position = end
        if kind == b"IEND":
            if length or position != len(data):
                raise ValueError("invalid PNG ending")
            return
    raise ValueError("PNG lacks IEND")


def decode_png(data):
    chunks = list(png_chunks(data))
    if not chunks or chunks[0][0] != b"IHDR" or len(chunks[0][1]) != 13:
        raise ValueError("invalid PNG header")
    width, height, depth, color, compression, filtering, interlace = struct.unpack(
        ">IIBBBBB", chunks[0][1]
    )
    channels = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}.get(color)
    valid_depths = {0: (1, 2, 4, 8, 16), 2: (8, 16), 3: (1, 2, 4, 8), 4: (8, 16), 6: (8, 16)}
    if (
        not channels
        or depth not in valid_depths[color]
        or compression
        or filtering
        or interlace
    ):
        raise ValueError("unsupported PNG layout (requires non-interlaced PNG)")
    if not width or not height or width * height > 16_000_000:
        raise ValueError("PNG dimensions exceed safety limit")
    stride = (width * channels * depth + 7) // 8
    expected = height * (stride + 1)
    if expected > 64 * 1024 * 1024:
        raise ValueError("decoded PNG exceeds safety limit")
    packed = b"".join(content for kind, content in chunks if kind == b"IDAT")
    decoder = zlib.decompressobj()
    raw = decoder.decompress(packed, expected + 1)
    if len(raw) != expected or not decoder.eof or decoder.unused_data:
        raise ValueError("PNG decompressed length mismatch")
    palette = next((content for kind, content in chunks if kind == b"PLTE"), b"")
    if color == 3 and (not palette or len(palette) % 3):
        raise ValueError("indexed PNG lacks a valid palette")
    bytes_per_pixel = max(1, (channels * depth + 7) // 8)
    rows, previous = [], bytearray(stride)
    for y in range(height):
        offset = y * (stride + 1)
        filter_type = raw[offset]
        row = bytearray(raw[offset + 1 : offset + 1 + stride])
        if filter_type > 4:
            raise ValueError("invalid PNG filter")
        if filter_type:
            for index in range(stride):
                left = row[index - bytes_per_pixel] if index >= bytes_per_pixel else 0
                above = previous[index]
                upper_left = previous[index - bytes_per_pixel] if index >= bytes_per_pixel else 0
                if filter_type == 1:
                    predictor = left
                elif filter_type == 2:
                    predictor = above
                elif filter_type == 3:
                    predictor = (left + above) // 2
                else:
                    estimate = left + above - upper_left
                    distances = abs(estimate - left), abs(estimate - above), abs(estimate - upper_left)
                    predictor = left if distances[0] <= distances[1] and distances[0] <= distances[2] else above if distances[1] <= distances[2] else upper_left
                row[index] = (row[index] + predictor) & 255
        rows.append(row)
        previous = row
    return width, height, depth, color, channels, palette, rows


def sample_rgb(decoded, edge):
    width, height, depth, color, channels, palette, rows = decoded
    scale = min(1, edge / max(width, height))
    target_width, target_height = max(1, int(width * scale)), max(1, int(height * scale))
    pixels = bytearray()
    for y in range(target_height):
        row = rows[min(height - 1, y * height // target_height)]
        for x in range(target_width):
            source_x = min(width - 1, x * width // target_width)
            if depth < 8:
                bit = source_x * depth
                value = (row[bit // 8] >> (8 - depth - bit % 8)) & ((1 << depth) - 1)
                values = [value if color == 3 else value * 255 // ((1 << depth) - 1)]
            else:
                step = depth // 8
                offset = source_x * channels * step
                values = [row[offset + channel * step] for channel in range(channels)]
            if color == 3:
                index = values[0] * 3
                if index + 3 > len(palette):
                    raise ValueError("PNG palette index out of range")
                rgb = palette[index : index + 3]
            elif color in (0, 4):
                rgb = [values[0]] * 3
            else:
                rgb = values[:3]
            if color in (4, 6):
                alpha = values[-1]
                rgb = [value * alpha // 255 for value in rgb]
            pixels.extend(rgb)
    return bytes(pixels), target_width, target_height


def preview_png(sample, edge):
    rgb, width, height = sample
    scale = min(1, edge / max(width, height))
    target_width, target_height = max(1, int(width * scale)), max(1, int(height * scale))
    raw = bytearray()
    for y in range(target_height):
        raw.append(0)
        for x in range(target_width):
            offset = (min(height - 1, y * height // target_height) * width + min(width - 1, x * width // target_width)) * 3
            raw.extend(value & 0xF8 for value in rgb[offset : offset + 3])

    def chunk(kind, content):
        return struct.pack(">I", len(content)) + kind + content + struct.pack(">I", zlib.crc32(kind + content) & 0xFFFFFFFF)

    data = PNG + chunk(b"IHDR", struct.pack(">IIBBBBB", target_width, target_height, 8, 2, 0, 0, 0)) + chunk(b"IDAT", zlib.compress(bytes(raw), 9)) + chunk(b"IEND", b"")
    return data, target_width, target_height


def progress(args, stage, **fields):
    print(json.dumps({"stage": stage, "elapsed_seconds": round(time.monotonic() - args.started, 2), **fields}, separators=(",", ":")), flush=True)


class Client:
    def __init__(self, base, deadline):
        self.base = base.rstrip("/")
        self.deadline = deadline

    def request(self, path, payload=None, cap=1024 * 1024):
        remaining = self.deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError("client deadline exceeded")
        data = None if payload is None else json.dumps(payload).encode()
        request = urllib.request.Request(self.base + path, data=data, headers={} if data is None else {"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=min(20, remaining)) as response:
            body = response.read(cap + 1)
            if len(body) > cap:
                raise ValueError("HTTP response exceeds byte limit")
            return body


def run(args, report, client):
    progress(args, "health")
    health = json.loads(client.request("/health"))
    if health.get("status") != "healthy":
        raise ValueError("telemetry health check failed")
    report["checks"]["health"] = True
    progress(args, "list")
    listing = json.loads(client.request("/api/images/list"))
    if listing.get("success") is not True or not isinstance(listing.get("data"), list):
        raise ValueError("invalid image-list response")
    names = sorted({name for name in listing["data"] if safe_name(name)})
    selected = [args.image] if names else []
    report.update(available_png_count=len(names), requested_count=1, requested=selected[:], selection="preferred filename with one fetch fallback", preferred_image=args.image, fallback_used=False)
    if not selected:
        raise ValueError("no PNG images are available")
    progress(args, "selected", available=len(names), count=1)
    name = selected[0]
    progress(args, "fetch", image=name, index=1, total=1)
    try:
        data = client.request("/api/images", {"images": [name]}, cap=MAX_IMAGE)
    except (OSError, ValueError) as error:
        fallback = next((candidate for candidate in names if candidate != name), None)
        if fallback is None or time.monotonic() >= client.deadline:
            raise
        report["fallback_used"] = True
        report["preferred_fetch_error"] = (type(error).__name__ + ": " + str(error))[:200]
        name = fallback
        report["requested"].append(name)
        progress(args, "fallback", image=fallback, preferred=args.image, reason=report["preferred_fetch_error"])
        data = client.request("/api/images", {"images": [name]}, cap=MAX_IMAGE)
    decoded = decode_png(data)
    sample = sample_rgb(decoded, 32)
    thumbnail, width, height = preview_png(sample, 32)
    metadata = {"name": name, "size_bytes": len(data), "sha256": hashlib.sha256(data).hexdigest(), "width": decoded[0], "height": decoded[1], "thumbnail": {"name": "thumbnails/01.png", "width": width, "height": height, "size_bytes": len(thumbnail), "sha256": hashlib.sha256(thumbnail).hexdigest()}}
    report["images"].append(metadata)
    report.update(status="passed", preview_method="nearest-neighbor RGB, 5 bits/channel; explicit alpha flattened to black", elapsed_seconds=round(time.monotonic() - args.started, 2))
    progress(args, "thumbnail_ready", images=1, max_edge=32, thumbnail_bytes=len(thumbnail))
    return thumbnail


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--api-base", default="http://satellite-telemetry.dphi-tm")
    parser.add_argument("--image", default="167.png")
    parser.add_argument("--max-seconds", type=int, default=300)
    parser.add_argument(
        "--report-output",
        type=Path,
        default=Path("/data/fisheye-results.json"),
    )
    parser.add_argument(
        "--thumbnail-output",
        type=Path,
        default=Path("/data/fisheye-thumbnail.png"),
    )
    args = parser.parse_args(argv)
    if not safe_name(args.image):
        parser.error("image must be a safe PNG filename, not a path")
    if not 1 <= args.max_seconds <= 3600:
        parser.error("max-seconds must be 1..3600")
    args.started = time.monotonic()
    report = {"version": 3, "status": "failed", "started_utc": datetime.now(timezone.utc).isoformat(), "images": [], "checks": {}}
    signal.signal(signal.SIGALRM, lambda *_: (_ for _ in ()).throw(TimeoutError("client deadline exceeded")))
    signal.alarm(args.max_seconds)
    try:
        args.report_output.parent.mkdir(parents=True, exist_ok=True)
        args.thumbnail_output.parent.mkdir(parents=True, exist_ok=True)
        args.thumbnail_output.unlink(missing_ok=True)
        thumbnail = run(args, report, Client(args.api_base, args.started + args.max_seconds))
    except Exception as error:
        report.update(status="failed", error=(type(error).__name__ + ": " + str(error))[:300], elapsed_seconds=round(time.monotonic() - args.started, 2))
        for metadata in report["images"]:
            metadata.pop("thumbnail", None)
        thumbnail = None
    finally:
        signal.alarm(0)
    report_temporary = args.report_output.with_suffix(".tmp")
    report_temporary.write_text(
        json.dumps(report, separators=(",", ":"), ensure_ascii=True), encoding="utf-8"
    )
    os.replace(report_temporary, args.report_output)
    if thumbnail is not None:
        thumbnail_temporary = args.thumbnail_output.with_suffix(".tmp")
        thumbnail_temporary.write_bytes(thumbnail)
        os.replace(thumbnail_temporary, args.thumbnail_output)
    print(json.dumps({"status": report["status"], "images": len(report.get("images", [])), "report_path": str(args.report_output), "thumbnail_path": str(args.thumbnail_output) if thumbnail is not None else None, "error": report.get("error")}, separators=(",", ":")), flush=True)
    return 0 if report["status"] == "passed" else 1


if __name__ == "__main__":
    sys.exit(main())