"""Image and video upscaling 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 (
    UPSCALE_POST_ENDPOINTS,
    UPSCALE_STATUS_ENDPOINTS,
    VIDEO_UPSCALE_POST_ENDPOINTS,
    VIDEO_UPSCALE_STATUS_ENDPOINT,
    UpscaleMode,
    VideoUpscaleMode,
    get_output_urls,
    get_status,
    get_task_id,
)


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

    # ------------------------------------------------------------------
    # Image upscaling
    # ------------------------------------------------------------------

    def upscale_image(
        self,
        mode: UpscaleMode,
        image_b64: str,
        scale_factor: Optional[str] = None,
        creativity: Optional[int] = None,
        prompt: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> str:
        """Submit an image upscale task. Returns task_id."""
        payload: dict[str, Any] = {"image": image_b64}

        if scale_factor:
            # Convert "2x" → 2, "4x" → 4
            factor = scale_factor.rstrip("xX")
            try:
                payload["scale_factor"] = int(factor)
            except ValueError:
                payload["scale_factor"] = 2

        if mode == UpscaleMode.CREATIVE:
            if creativity is not None:
                payload["creativity"] = creativity
            if prompt:
                payload["prompt"] = prompt

        if seed is not None:
            payload["seed"] = seed

        endpoint = UPSCALE_POST_ENDPOINTS[mode]
        raw = self._client.post(endpoint, json=payload)
        return get_task_id(raw)

    def upscale_image_status(self, mode: UpscaleMode, task_id: str) -> Tuple[str, dict[str, Any]]:
        endpoint = UPSCALE_STATUS_ENDPOINTS[mode].format(task_id=task_id)
        raw = self._client.get(endpoint)
        return get_status(raw), raw

    # ------------------------------------------------------------------
    # Video upscaling
    # ------------------------------------------------------------------

    def upscale_video(
        self,
        mode: VideoUpscaleMode,
        video_b64: str,
    ) -> str:
        """Submit a video upscale task. Returns task_id."""
        payload: dict[str, Any] = {"video": video_b64}
        endpoint = VIDEO_UPSCALE_POST_ENDPOINTS[mode]
        raw = self._client.post(endpoint, json=payload)
        return get_task_id(raw)

    def upscale_video_status(self, task_id: str) -> Tuple[str, dict[str, Any]]:
        endpoint = VIDEO_UPSCALE_STATUS_ENDPOINT.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)