"""Image tensor output helpers."""

from __future__ import annotations

from io import BytesIO
from pathlib import Path
import subprocess
import tempfile

import numpy as np
from PIL import Image
import torch

from diffusion_cli.config import (
    DEFAULT_OUTPUT_EXTENSION,
    DEFAULT_OUTPUT_QUALITY,
)
from diffusion_cli.errors import DiffusionCliError


def outputPathWithExtension(output: Path, extension: str) -> Path:
    """Return output with the normalized image extension."""

    return output.with_suffix(f".{extension}")


def outputPaths(
    output: Path,
    batch_size: int,
    output_extension: str = DEFAULT_OUTPUT_EXTENSION,
) -> list[Path]:
    """Return output file names for a batch."""

    output = outputPathWithExtension(output, output_extension)
    if batch_size == 1:
        return [output]
    return [
        output.with_name(f"{output.stem}_{index:04d}{output.suffix}")
        for index in range(batch_size)
    ]


def saveImages(
    images: torch.Tensor,
    output: Path,
    output_extension: str = DEFAULT_OUTPUT_EXTENSION,
    output_quality: int = DEFAULT_OUTPUT_QUALITY,
) -> list[Path]:
    """Save an NCHW image tensor in [0, 1] as final image files."""

    image_arrays = imageArrays(images)
    paths = outputPaths(output, len(image_arrays), output_extension)

    for image_array, path in zip(image_arrays, paths, strict=True):
        if output_extension == "png":
            Image.fromarray(np.asarray(image_array), mode="RGB").save(path)
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_png = Path(temp_dir) / "input.png"
                Image.fromarray(np.asarray(image_array), mode="RGB").save(
                    temp_png,
                )
                convertPngFile(
                    temp_png,
                    path,
                    output_extension,
                    output_quality,
                )

    return paths


def imageArrays(images: torch.Tensor) -> np.ndarray:
    """Convert an NCHW image tensor in [0, 1] to NHWC uint8 arrays."""

    if images.ndim != 4:
        raise ValueError(f"Expected NCHW image tensor, got {images.shape}")
    if images.shape[1] != 3:
        raise ValueError(f"Expected three image channels, got {images.shape}")

    return (
        images.detach()
        .float()
        .clamp(0, 1)
        .permute(0, 2, 3, 1)
        .mul(255)
        .round()
        .to(torch.uint8)
        .cpu()
        .numpy()
    )


def encodeImages(
    images: torch.Tensor,
    output_extension: str = DEFAULT_OUTPUT_EXTENSION,
    output_quality: int = DEFAULT_OUTPUT_QUALITY,
) -> list[bytes]:
    """Encode an NCHW image tensor in [0, 1] to final image bytes."""

    image_arrays = imageArrays(images)
    return [
        encodeArray(image_array, output_extension, output_quality)
        for image_array in image_arrays
    ]


def convertPngFile(
    png_path: Path,
    output_path: Path,
    output_extension: str,
    output_quality: int,
) -> None:
    """Convert one PNG file to the requested output format."""

    try:
        subprocess.run(
            [
                "magick",
                str(png_path),
                "-quality",
                str(output_quality),
                str(output_path),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    except FileNotFoundError as exc:
        raise DiffusionCliError(
            "ImageMagick executable not found: magick"
        ) from exc
    except subprocess.CalledProcessError as exc:
        message = (exc.stderr or "").strip() or (exc.stdout or "").strip()
        if not message:
            message = f"exit code {exc.returncode}"
        raise DiffusionCliError(
            f"ImageMagick failed to write {output_extension} output: "
            f"{message}"
        ) from exc

    if not output_path.is_file():
        raise DiffusionCliError(
            f"ImageMagick did not create output file: {output_path}"
        )


def encodeArray(
    image_array: np.ndarray,
    output_extension: str,
    output_quality: int,
) -> bytes:
    """Encode one RGB image array to final image bytes."""

    if output_extension == "png":
        output = BytesIO()
        Image.fromarray(np.asarray(image_array), mode="RGB").save(
            output,
            format="PNG",
        )
        return output.getvalue()

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)
        input_png = temp_path / "input.png"
        output_path = temp_path / f"output.{output_extension}"
        Image.fromarray(np.asarray(image_array), mode="RGB").save(input_png)
        convertPngFile(
            input_png,
            output_path,
            output_extension,
            output_quality,
        )
        return output_path.read_bytes()