import time
from typing import Any, Dict, Optional

from ..base_client import BaseClient
from ..models import GenerationResponse, TaskStatus


class Image:
    """Image generation resource."""

    def __init__(self, client: BaseClient):
        self._client = client

    def generate(self, **kwargs) -> Dict[str, Any]:
        """Synchronous image generation. Returns result directly."""
        return self._client.post("/v1/images/generations", kwargs)

    def generate_async(self, **kwargs) -> GenerationResponse:
        """Submit an async image generation task."""
        data = self._client.post("/v1/images/generations", kwargs)
        return GenerationResponse(data)

    def query_task(self, task_id: str) -> TaskStatus:
        """Query the status of an image generation task."""
        data = self._client.get(f"/{task_id}")
        return TaskStatus(data)

    def run(
        self,
        options: dict,
        poll_interval: float = 2.0,
        timeout: float = 300.0,
    ) -> TaskStatus:
        """Submit and poll until completion."""
        response = self.generate_async(**options)

        if not response.task_id:
            raise RuntimeError("No task ID returned from generation request")

        return self._poll(response.task_id, poll_interval, timeout)

    def _poll(self, task_id: str, poll_interval: float, timeout: float) -> TaskStatus:
        start = time.time()
        while True:
            status = self.query_task(task_id)
            if status.is_completed() or status.is_failed():
                return status
            if time.time() - start > timeout:
                raise TimeoutError(f"Task {task_id} timed out after {timeout}s")
            time.sleep(poll_interval)
