import base64

import httpx

from spaik_sdk.image_gen.options import ImageGenOptions

GOOGLE_GENERATIVE_ENDPOINT = "https://generativelanguage.googleapis.com/v1/models"


def _derive_resolution_and_aspect(width: int, height: int) -> tuple[str, str]:
    from math import gcd

    divisor = gcd(width, height)
    aspect_w = width // divisor
    aspect_h = height // divisor
    aspect_ratio = f"{aspect_w}:{aspect_h}"

    max_dim = max(width, height)
    if max_dim <= 1024:
        resolution = "1K"
    elif max_dim <= 2048:
        resolution = "2K"
    else:
        resolution = "4K"

    return resolution, aspect_ratio


async def generate_image(
    prompt: str,
    model: str,
    api_key: str,
    options: ImageGenOptions,
    endpoint: str | None = None,
    headers: dict[str, str] | None = None,
) -> bytes:
    base_url = endpoint or GOOGLE_GENERATIVE_ENDPOINT
    url = f"{base_url}/{model}:generateContent?key={api_key}"

    request_headers = {
        "Content-Type": "application/json",
    }
    if headers:
        request_headers.update(headers)

    resolution, aspect_ratio = _derive_resolution_and_aspect(options.width, options.height)

    generation_config: dict = {
        "responseModalities": ["image", "text"],
        "resolution": resolution,
        "aspectRatio": aspect_ratio,
    }
    generation_config.update(options.vendor)

    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": generation_config,
    }

    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(url, headers=request_headers, json=payload)
        response.raise_for_status()
        data = response.json()

    candidates = data.get("candidates", [])
    if not candidates:
        raise ValueError("No image generated by Gemini API")

    parts = candidates[0].get("content", {}).get("parts", [])
    for part in parts:
        if "inlineData" in part:
            image_b64 = part["inlineData"]["data"]
            return base64.b64decode(image_b64)

    raise ValueError("No image data found in Gemini API response")
