"""Image generation and analysis API methods."""

from __future__ import annotations

from typing import Any, Optional, Tuple

from freepik_cli.api.client import FreepikClient
from freepik_cli.api.models import (
    IMAGE_POST_ENDPOINTS,
    IMAGE_STATUS_ENDPOINTS,
    ImageModel,
    get_output_urls,
    get_status,
    get_task_id,
)


class ImageAPI:
    def __init__(self, client: FreepikClient) -> None:
        self._client = client

    def generate(self, model: ImageModel, payload: dict[str, Any]) -> str:
        """Submit a generation task. Returns task_id."""
        endpoint = IMAGE_POST_ENDPOINTS[model]
        raw = self._client.post(endpoint, json=payload)
        return get_task_id(raw)

    def get_status(self, model: ImageModel, task_id: str) -> Tuple[str, dict[str, Any]]:
        """Poll status. Returns (status_str, raw_response)."""
        endpoint = IMAGE_STATUS_ENDPOINTS[model].format(task_id=task_id)
        raw = self._client.get(endpoint)
        return get_status(raw), raw

    def get_output_urls(self, raw: dict[str, Any]) -> list[str]:
        return get_output_urls(raw)

    # ------------------------------------------------------------------
    # Image-to-prompt (describe)
    # ------------------------------------------------------------------

    def describe_submit(self, image_b64: str) -> str:
        """Submit image-to-prompt task. Returns task_id."""
        raw = self._client.post("/v1/ai/image-to-prompt", json={"image": image_b64})
        return get_task_id(raw)

    def describe_status(self, task_id: str) -> Tuple[str, dict[str, Any]]:
        raw = self._client.get(f"/v1/ai/image-to-prompt/{task_id}")
        return get_status(raw), raw

    def get_prompt_text(self, raw: dict[str, Any]) -> str:
        """Extract generated prompt text from a completed describe response."""
        data = raw.get("data", raw)
        return (
            data.get("prompt")
            or data.get("description")
            or data.get("text")
            or data.get("result", {}).get("prompt", "")
            or ""
        )

    # ------------------------------------------------------------------
    # Image expansion (outpainting)
    # ------------------------------------------------------------------

    def expand_submit(
        self,
        model: str,
        image_b64: str,
        left: int = 0,
        right: int = 0,
        top: int = 0,
        bottom: int = 0,
        prompt: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> str:
        payload: dict[str, Any] = {
            "image": image_b64,
            "left": left,
            "right": right,
            "top": top,
            "bottom": bottom,
        }
        if prompt:
            payload["prompt"] = prompt
        if seed is not None:
            payload["seed"] = seed

        endpoint_map = {
            "flux-pro": "/v1/ai/image-expand/flux-pro",
            "ideogram": "/v1/ai/image-expand/ideogram",
            "seedream-v4-5": "/v1/ai/image-expand/seedream-v4-5",
        }
        endpoint = endpoint_map.get(model, "/v1/ai/image-expand/flux-pro")
        raw = self._client.post(endpoint, json=payload)
        return get_task_id(raw)

    def expand_status(self, model: str, task_id: str) -> Tuple[str, dict[str, Any]]:
        endpoint_map = {
            "flux-pro": "/v1/ai/image-expand/flux-pro",
            "ideogram": "/v1/ai/image-expand/ideogram",
            "seedream-v4-5": "/v1/ai/image-expand/seedream-v4-5",
        }
        base = endpoint_map.get(model, "/v1/ai/image-expand/flux-pro")
        raw = self._client.get(f"{base}/{task_id}")
        return get_status(raw), raw