#!/usr/bin/env python3
"""Generate a PNG through the default PinAI Images API."""

import argparse
import base64
import json
import os
import re
import struct
import sys
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Mapping, Optional
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import Request, urlopen

DEFAULT_API_URL = "https://api.pinaic.com/v1"
DEFAULT_MODEL = "gpt-image-2"
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
PNG_COLOR_TYPES = {0: "grayscale", 2: "RGB", 3: "indexed color", 4: "grayscale with alpha", 6: "RGBA"}


@dataclass(frozen=True)
class ImageGeneration:
    """Generated image data plus client-side timing details."""

    image: bytes
    api_request_seconds: float
    image_download_seconds: float
    response_mode: str


def extract_image_bytes(
    payload: Mapping[str, Any],
    download_url: Optional[Callable[[str], bytes]] = None,
) -> bytes:
    """Return the first image as bytes from a standard Images API response."""
    data = payload.get("data")
    if not isinstance(data, list) or not data or not isinstance(data[0], Mapping):
        raise ValueError("API response does not contain data[0]")

    image = data[0]
    encoded = image.get("b64_json")
    if isinstance(encoded, str) and encoded:
        try:
            return base64.b64decode(encoded, validate=True)
        except (ValueError, TypeError) as error:
            raise ValueError("API response contains invalid b64_json") from error

    image_url = image.get("url")
    if isinstance(image_url, str) and image_url:
        return (download_url or download_image)(image_url)

    raise ValueError("API response data[0] must contain b64_json or url")


def download_image(image_url: str, timeout: int = 120) -> bytes:
    parsed = urlsplit(image_url)
    if parsed.scheme != "https" or not parsed.netloc:
        raise ValueError("Image URL must use HTTPS")
    request = Request(image_url, headers={"User-Agent": "openai-imagegen-skill/1.0"})
    try:
        with urlopen(request, timeout=timeout) as response:
            return response.read()
    except (HTTPError, URLError, TimeoutError) as error:
        raise RuntimeError("Could not download image returned by provider") from error


def redact_message(message: str, api_key: str) -> str:
    cleaned = message.replace(api_key, "[REDACTED]") if api_key else message
    return re.sub(
        r"(?i)(authorization|api[_-]?key|token)(\s*[:=]\s*)([^\s,;]+)",
        r"\1\2[REDACTED]",
        cleaned,
    )


def images_endpoint(api_url: str) -> str:
    """Accept either an API base URL or a complete images/generations endpoint."""
    normalized = api_url.rstrip("/")
    if normalized.endswith("/images/generations"):
        return normalized
    return normalized + "/images/generations"


def request_image_with_timing(
    *,
    prompt: str,
    api_key: str,
    api_url: str,
    model: str,
    size: Optional[str],
    quality: Optional[str],
    timeout: int,
) -> ImageGeneration:
    body: dict[str, str] = {"model": model, "prompt": prompt, "output_format": "png"}
    if size:
        body["size"] = size
    if quality:
        body["quality"] = quality
    endpoint = images_endpoint(api_url)
    request = Request(
        endpoint,
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    api_started = time.perf_counter()
    try:
        with urlopen(request, timeout=timeout) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except HTTPError as error:
        details = error.read().decode("utf-8", errors="replace")[:1000]
        raise RuntimeError(
            f"Images API returned HTTP {error.code}: {redact_message(details, api_key)}"
        ) from error
    except (URLError, TimeoutError) as error:
        raise RuntimeError("Could not reach the Images API") from error
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise RuntimeError("Images API returned invalid JSON") from error
    api_request_seconds = time.perf_counter() - api_started

    data = payload.get("data")
    first_image = data[0] if isinstance(data, list) and data else {}
    response_mode = "b64_json"
    if isinstance(first_image, Mapping) and not first_image.get("b64_json") and first_image.get("url"):
        response_mode = "url"
    download_started = time.perf_counter()
    image = extract_image_bytes(payload, lambda url: download_image(url, timeout))
    image_download_seconds = time.perf_counter() - download_started if response_mode == "url" else 0.0
    return ImageGeneration(image, api_request_seconds, image_download_seconds, response_mode)


def request_image(
    *, prompt: str, api_key: str, base_url: str, model: str, size: Optional[str], timeout: int
) -> bytes:
    """Return generated bytes for callers that do not need timing details."""
    return request_image_with_timing(
        prompt=prompt,
        api_key=api_key,
        api_url=base_url,
        model=model,
        size=size,
        quality=None,
        timeout=timeout,
    ).image


def save_png(output: Path, image: bytes, force: bool) -> None:
    if not image.startswith(PNG_SIGNATURE):
        raise ValueError("Provider response is not PNG data")
    output.parent.mkdir(parents=True, exist_ok=True)
    if output.is_symlink():
        raise ValueError("Refusing to write through a symbolic-link output path")
    if output.exists() and not force:
        raise FileExistsError(f"Output already exists: {output}. Use --force to replace it.")
    with tempfile.NamedTemporaryFile(dir=output.parent, delete=False) as handle:
        temporary_path = Path(handle.name)
        handle.write(image)
    try:
        os.replace(temporary_path, output)
    finally:
        temporary_path.unlink(missing_ok=True)


def png_metadata(image: bytes) -> tuple[int, int, int, str]:
    """Return width, height, bit depth, and color mode from a PNG IHDR chunk."""
    if not image.startswith(PNG_SIGNATURE) or len(image) < 29 or image[12:16] != b"IHDR":
        raise ValueError("Provider response is not a valid PNG with an IHDR chunk")
    width, height, bit_depth, color_type = struct.unpack(">IIBB", image[16:26])
    color_mode = PNG_COLOR_TYPES.get(color_type, f"unknown color type {color_type}")
    return width, height, bit_depth, color_mode


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate a PNG through the default PinAI Images API.")
    parser.add_argument("--prompt", required=True, help="Image description")
    parser.add_argument("--output", required=True, type=Path, help="Destination .png file")
    parser.add_argument("--size", help="Optional image size, for example 1024x1024")
    parser.add_argument(
        "--quality",
        default=os.getenv("IMAGE_API_QUALITY") or os.getenv("PINAI_IMAGE_QUALITY"),
        help="Optional provider quality setting, for example low, medium, high, or auto",
    )
    parser.add_argument(
        "--api-url",
        "--base-url",
        dest="api_url",
        default=os.getenv("IMAGE_API_URL") or os.getenv("PINAI_BASE_URL") or DEFAULT_API_URL,
        help="Optional override for the default PinAI API base URL or complete endpoint",
    )
    parser.add_argument(
        "--model",
        default=os.getenv("IMAGE_API_MODEL") or os.getenv("PINAI_IMAGE_MODEL") or DEFAULT_MODEL,
        help=f"Image model (default: {DEFAULT_MODEL})",
    )
    parser.add_argument("--timeout", type=int, default=120, help="Network timeout in seconds")
    parser.add_argument("--force", action="store_true", help="Replace an existing output file")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    api_key = os.getenv("IMAGE_API_KEY") or os.getenv("PINAI_API_KEY")
    if not api_key:
        print("PINAI_API_KEY is required in the environment.", file=sys.stderr)
        return 2
    if not args.api_url:
        print("PinAI API URL is not configured.", file=sys.stderr)
        return 2
    if args.timeout <= 0:
        print("--timeout must be a positive integer.", file=sys.stderr)
        return 2
    if args.output.suffix.lower() != ".png":
        print("--output must use a .png extension.", file=sys.stderr)
        return 2
    started_at = datetime.now(timezone.utc)
    total_started = time.perf_counter()
    try:
        generation = request_image_with_timing(
            prompt=args.prompt,
            api_key=api_key,
            api_url=args.api_url,
            model=args.model,
            size=args.size,
            quality=args.quality,
            timeout=args.timeout,
        )
        width, height, bit_depth, color_mode = png_metadata(generation.image)
        save_started = time.perf_counter()
        save_png(args.output, generation.image, args.force)
        save_seconds = time.perf_counter() - save_started
    except (FileExistsError, RuntimeError, ValueError) as error:
        print(redact_message(str(error), api_key), file=sys.stderr)
        return 1
    total_seconds = time.perf_counter() - total_started
    finished_at = datetime.now(timezone.utc)
    print(f"Saved PNG: {args.output}")
    print(f"Started at (UTC): {started_at.isoformat(timespec='seconds')}")
    print(f"Finished at (UTC): {finished_at.isoformat(timespec='seconds')}")
    print(f"Provider request: {generation.api_request_seconds:.2f}s")
    print(f"Image download: {generation.image_download_seconds:.2f}s")
    print(f"File save: {save_seconds:.2f}s")
    print(f"Total elapsed: {total_seconds:.2f}s")
    print(f"Response mode: {generation.response_mode}")
    print(f"Actual dimensions: {width}x{height}px")
    print(f"Requested size: {args.size or 'provider default'}")
    print(f"Requested quality: {args.quality or 'provider default'}")
    print(f"Image quality: lossless PNG, {bit_depth}-bit {color_mode}")
    file_size = args.output.stat().st_size
    print(f"File size: {file_size} bytes ({file_size / (1024 * 1024):.2f} MiB)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())