"""Seedream image client for O1Key's asynchronous image API."""

from __future__ import annotations

import os
import time
from typing import Any, Awaitable, Callable, Optional, Sequence

from PIL import Image

from ..utils.nano_banana_async import (
    extract_async_image_result_urls,
    image_to_upload_payload,
    parse_completed_async_image_task,
    poll_async_image_task,
    submit_async_image_task,
    upload_images_to_temp_urls,
)
from ..utils.o1key_image_catalog import (
    MAX_UNIFIED_REFERENCE_IMAGES,
    SEEDREAM_MODEL_OPTIONS,
    SEEDREAM_LAYER_RESOLUTION_OPTIONS,
    SEEDREAM_OUTPUT_FORMAT_OPTIONS,
    SEEDREAM_SIZE_MATRIX,
    UNIFIED_IMAGE_ROUTE_OPTIONS,
)


SEEDREAM_API_MODEL_ID = "dola-seedream-5-0-pro-260628-ep"
SEEDREAM_REFERENCE_MAX_BYTES = 30 * 1024 * 1024
SEEDREAM_REFERENCE_MAX_PIXELS = 6000 * 6000
SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE = 14
SEEDREAM_REFERENCE_MIN_ASPECT_RATIO = 1 / 16
SEEDREAM_REFERENCE_MAX_ASPECT_RATIO = 16
SEEDREAM_LAYER_REFERENCE_MIN_PIXELS = 512 * 512


def validate_seedream_reference_dimensions(
    width: int,
    height: int,
    *,
    label: str = "Seedream 参考图",
    layer_decomposition: bool = False,
) -> None:
    """Validate the current Volcengine per-image reference-size contract."""
    width = int(width)
    height = int(height)
    if width <= 0 or height <= 0:
        raise ValueError(f"{label}尺寸无效：{width}×{height}")

    pixels = width * height
    if (
        not layer_decomposition
        and (
            width <= SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE
            or height <= SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE
        )
    ):
        raise ValueError(
            f"{label}宽和高都必须大于 {SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE}px，"
            f"当前为 {width}×{height}"
        )
    ratio = width / height
    if (
        ratio < SEEDREAM_REFERENCE_MIN_ASPECT_RATIO
        or ratio > SEEDREAM_REFERENCE_MAX_ASPECT_RATIO
    ):
        raise ValueError(
            f"{label}宽高比必须在 1:16～16:1，当前为 {width}:{height}"
        )
    if layer_decomposition:
        if not (
            SEEDREAM_LAYER_REFERENCE_MIN_PIXELS
            <= pixels
            <= SEEDREAM_REFERENCE_MAX_PIXELS
        ):
            raise ValueError(
                f"{label}总像素必须在 512×512（262144）～6000×6000（36000000）之间，"
                f"当前为 {width}×{height}（{pixels}）"
            )
        return
    if pixels > SEEDREAM_REFERENCE_MAX_PIXELS:
        raise ValueError(
            f"{label}总像素不能超过 6000×6000（36000000），"
            f"当前为 {width}×{height}（{pixels}）"
        )


def validate_seedream_reference_image(
    image: Image.Image,
    *,
    label: str = "Seedream 参考图",
    layer_decomposition: bool = False,
) -> None:
    """Validate reference dimensions and the exact bytes sent to the uploader."""
    validate_seedream_reference_dimensions(
        image.width,
        image.height,
        label=label,
        layer_decomposition=layer_decomposition,
    )
    payload, _extension, _content_type = image_to_upload_payload(image)
    try:
        payload_size = (
            os.path.getsize(payload)
            if isinstance(payload, (str, os.PathLike))
            else len(payload)
        )
    except (OSError, TypeError) as exc:
        raise ValueError(f"无法读取{label}文件大小") from exc
    if payload_size > SEEDREAM_REFERENCE_MAX_BYTES:
        raise ValueError(
            f"{label}文件不能超过 30MB，当前为 {payload_size / 1024 / 1024:.2f}MB"
        )


def validate_seedream_reference_images(
    images: Sequence[Image.Image],
    *,
    layer_decomposition: bool = False,
) -> None:
    for index, image in enumerate(images, start=1):
        validate_seedream_reference_image(
            image,
            label=f"Seedream 参考图{index}",
            layer_decomposition=layer_decomposition,
        )


def resolve_seedream_model(model_name: str, route: str) -> str:
    """Map the stable workflow value to Seedream's API model identifier."""
    if model_name == SEEDREAM_API_MODEL_ID:
        return model_name
    if model_name not in SEEDREAM_MODEL_OPTIONS:
        raise ValueError(f"Seedream 模型无效：{model_name}")
    if route not in UNIFIED_IMAGE_ROUTE_OPTIONS:
        raise ValueError(f"Seedream 模型线路无效：{route}")
    # O1Key currently exposes one Seedream endpoint for every displayed route.
    return SEEDREAM_API_MODEL_ID


def build_seedream_submit_body(
    *,
    model: str,
    prompt: str,
    size: Optional[str],
    output_format: str,
    image_urls: Optional[Sequence[str]] = None,
    layer_decomposition: bool = False,
) -> dict[str, Any]:
    """Build and validate the paid Seedream request without logging URLs."""
    normalized_prompt = str(prompt or "").strip()
    if not normalized_prompt and not layer_decomposition:
        raise ValueError("请输入提示词")
    if model != SEEDREAM_API_MODEL_ID:
        raise ValueError(f"Seedream API 模型无效：{model}")

    normalized_format = str(output_format or "").strip().lower()
    if normalized_format not in SEEDREAM_OUTPUT_FORMAT_OPTIONS:
        raise ValueError("Seedream 输出格式仅支持 png 或 jpeg")

    if layer_decomposition and normalized_format != "png":
        raise ValueError("Seedream 图层拆分仅支持 png 输出格式")

    normalized_size = str(size or "").strip().lower().replace("*", "x").replace("×", "x")
    if normalized_size:
        if layer_decomposition:
            normalized_size = "auto" if normalized_size == "auto" else normalized_size.upper()
            if normalized_size not in SEEDREAM_LAYER_RESOLUTION_OPTIONS:
                raise ValueError(f"Seedream 图层拆分分辨率无效：{size}")
        elif normalized_size not in set(SEEDREAM_SIZE_MATRIX.values()):
            raise ValueError(f"Seedream 图片尺寸无效：{size}")

    urls = [str(url or "").strip() for url in (image_urls or ())]
    if len(urls) > MAX_UNIFIED_REFERENCE_IMAGES:
        raise ValueError(f"Seedream 参考图最多支持 {MAX_UNIFIED_REFERENCE_IMAGES} 张")
    if any(not url.startswith("https://") for url in urls):
        raise ValueError("Seedream 参考图必须使用临时素材 HTTPS URL")
    if layer_decomposition and len(urls) != 1:
        raise ValueError("Seedream 图层拆分必须且只能提供1张参考图")

    body: dict[str, Any] = {
        "model": model,
        "n": 1,
        "output_format": normalized_format,
        "watermark": False,
    }
    if normalized_size:
        body["size"] = normalized_size
    if normalized_prompt:
        body["prompt"] = normalized_prompt
    if urls:
        body["images"] = urls
    if layer_decomposition:
        body["layer_decomposition"] = True
    return body


def _seedream_result_items(payload: Any) -> list[dict[str, Any]]:
    """Return the first documented image-item list without exposing its URLs."""
    pending = [payload]
    seen: set[int] = set()
    while pending:
        value = pending.pop(0)
        if not isinstance(value, dict) or id(value) in seen:
            continue
        seen.add(id(value))
        images = value.get("images")
        if isinstance(images, list) and all(isinstance(item, dict) for item in images):
            return images
        for key in ("data", "result", "output"):
            nested = value.get(key)
            if isinstance(nested, dict):
                pending.append(nested)
    return []


def _bounded_int_list(value: Any, *, length: int) -> list[int] | None:
    if not isinstance(value, (list, tuple)) or len(value) != length:
        return None
    try:
        return [int(item) for item in value]
    except (TypeError, ValueError):
        return None


def extract_seedream_layer_metadata(payload: Any) -> list[dict[str, Any]]:
    """Sanitize layer metadata; result URLs are deliberately excluded."""
    metadata: list[dict[str, Any]] = []
    for offset, item in enumerate(_seedream_result_items(payload)):
        try:
            z_index = max(0, min(16, int(item.get("z_index", offset))))
        except (TypeError, ValueError):
            z_index = offset
        safe: dict[str, Any] = {"z_index": z_index}
        for key, limit in (("name", 200), ("description", 1000), ("size", 64), ("output_format", 16)):
            value = item.get(key)
            if isinstance(value, str) and value.strip():
                safe[key] = value.strip()[:limit]
        bounding_box = item.get("bounding_box")
        if isinstance(bounding_box, dict):
            absolute = _bounded_int_list(bounding_box.get("absolute"), length=4)
            normalized = _bounded_int_list(bounding_box.get("normalized"), length=4)
            safe_box = {}
            if absolute is not None:
                safe_box["absolute"] = absolute
            if normalized is not None:
                safe_box["normalized"] = normalized
            if safe_box:
                safe["bounding_box"] = safe_box
        metadata.append(safe)
    return metadata


class SeedreamImageClient:
    """Upload references, submit one Seedream task, poll it, and decode results."""

    def __init__(self, *, base_url: str, api_key: str):
        self.base_url = str(base_url).rstrip("/")
        self.api_key = api_key

    async def generate_async(
        self,
        *,
        session: Any,
        prompt: str,
        model: str,
        size: Optional[str],
        output_format: str,
        images: Optional[Sequence[Image.Image]] = None,
        layer_decomposition: bool = False,
        upload_cache: Optional[dict[int, Awaitable[str]]] = None,
        check_interrupt: Optional[Callable[[], None]] = None,
        progress_callback: Optional[Callable[[float], None]] = None,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_task_success: bool = True,
        task_completed_callback: Optional[
            Callable[[str, int, float, list[str]], None]
        ] = None,
    ) -> tuple[list[Image.Image], dict[str, Any]]:
        if check_interrupt:
            check_interrupt()

        reference_images = list(images or ())
        validate_seedream_reference_images(
            reference_images,
            layer_decomposition=layer_decomposition,
        )

        task_started = time.time()
        image_urls = await upload_images_to_temp_urls(
            session=session,
            base_url=self.base_url,
            api_key=self.api_key,
            images=reference_images,
            node_label="Seedream",
            check_interrupt=check_interrupt,
            upload_cache=upload_cache,
            log_success=log_task_success,
        )
        body = build_seedream_submit_body(
            model=model,
            prompt=prompt,
            size=size,
            output_format=output_format,
            image_urls=image_urls,
            layer_decomposition=layer_decomposition,
        )

        task_id = await submit_async_image_task(
            session,
            self.base_url,
            self.api_key,
            body,
            "Seedream",
            log_body_enabled=False,
            log_success=log_task_success,
        )
        task_payload = await poll_async_image_task(
            session,
            self.base_url,
            self.api_key,
            task_id,
            "Seedream",
            check_interrupt=check_interrupt,
            log_body_enabled=False,
            progress_callback=progress_callback,
            log_success=log_task_success,
        )
        task_done = time.time()
        parse_started = time.time()
        task_payload, parsed = await parse_completed_async_image_task(
            task_payload,
            session,
            self.base_url,
            self.api_key,
            task_id,
            "Seedream",
            check_interrupt=check_interrupt,
            result_url_callback=None,
            log_downloads=log_downloads,
        )
        if isinstance(parsed, tuple) and len(parsed) == 2:
            result_images, metrics = parsed
        else:
            result_images = parsed
            metrics = {
                "download_bytes": 0,
                "download_seconds": 0.0,
                "download_wall_seconds": 0.0,
                "inline_images": 0,
            }

        result_urls = extract_async_image_result_urls(task_payload)
        result_metadata = extract_seedream_layer_metadata(task_payload)
        if layer_decomposition:
            paired = []
            for index, image in enumerate(result_images):
                metadata = (
                    result_metadata[index]
                    if index < len(result_metadata)
                    else {"z_index": index}
                )
                setattr(image, "_o1key_seedream_layer", metadata)
                paired.append((metadata.get("z_index", index), index, image))
            paired.sort(key=lambda item: (item[0], item[1]))
            result_images = [item[2] for item in paired]
            result_metadata = [
                getattr(image, "_o1key_seedream_layer", {"z_index": index})
                for index, image in enumerate(result_images)
            ]
        if result_url_callback:
            for url in result_urls:
                result_url_callback(url)
        if task_completed_callback:
            task_completed_callback(
                task_id,
                len(result_images),
                time.time() - task_started,
                result_urls,
            )

        return result_images, {
            "task_id": task_id,
            "task_ids": [task_id],
            "task_ms": (task_done - task_started) * 1000,
            "parse_ms": (time.time() - parse_started) * 1000,
            "download_ms": metrics["download_wall_seconds"] * 1000,
            "download_total_ms": metrics["download_seconds"] * 1000,
            "download_bytes": metrics["download_bytes"],
            "inline_images": metrics["inline_images"],
            "result_metadata": result_metadata,
        }


__all__ = [
    "SEEDREAM_API_MODEL_ID",
    "SEEDREAM_LAYER_REFERENCE_MIN_PIXELS",
    "SEEDREAM_REFERENCE_MAX_BYTES",
    "SEEDREAM_REFERENCE_MAX_PIXELS",
    "SeedreamImageClient",
    "build_seedream_submit_body",
    "extract_seedream_layer_metadata",
    "resolve_seedream_model",
    "validate_seedream_reference_dimensions",
    "validate_seedream_reference_image",
    "validate_seedream_reference_images",
]