import time
from typing import Any, Dict

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


class Video:
    """Video generation resource supporting Seedance and Kling."""

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

    # ---- Seedance ----

    def generate_async(self, **kwargs) -> GenerationResponse:
        """Submit a Seedance video generation task."""
        data = self._client.post("/v1/video/generations", kwargs)
        return GenerationResponse(data)

    def query_task(self, task_id: str) -> TaskStatus:
        """Query Seedance task status."""
        data = self._client.get(f"/v1/video/generations/{task_id}")
        return TaskStatus(data)

    def run(
        self,
        options: dict,
        poll_interval: float = 2.0,
        timeout: float = 300.0,
    ) -> TaskStatus:
        """Submit Seedance task 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(
            lambda tid: self.query_task(tid),
            response.task_id,
            poll_interval,
            timeout,
        )

    # ---- Kling ----

    def kling_text2video(self, **kwargs) -> GenerationResponse:
        """Submit a Kling text-to-video task."""
        data = self._client.post("/kling/v1/videos/text2video", kwargs)
        return GenerationResponse(data)

    def kling_image2video(self, **kwargs) -> GenerationResponse:
        """Submit a Kling image-to-video task."""
        data = self._client.post("/kling/v1/videos/image2video", kwargs)
        return GenerationResponse(data)

    def kling_multi_image2video(self, **kwargs) -> GenerationResponse:
        """Submit a Kling multi-image-to-video task."""
        data = self._client.post("/kling/v1/videos/multi-image2video", kwargs)
        return GenerationResponse(data)

    def kling_query_task(self, action: str, task_id: str) -> TaskStatus:
        """Query Kling task status by action and task ID."""
        data = self._client.get(f"/kling/v1/videos/{action}/{task_id}")
        return TaskStatus(data)

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

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

        return self._poll(
            lambda tid: self.kling_query_task("text2video", tid),
            response.task_id,
            poll_interval,
            timeout,
        )

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

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

        return self._poll(
            lambda tid: self.kling_query_task("image2video", tid),
            response.task_id,
            poll_interval,
            timeout,
        )

    def run_kling_multi_image2video(
        self,
        options: dict,
        poll_interval: float = 2.0,
        timeout: float = 300.0,
    ) -> TaskStatus:
        """Submit Kling multi-image2video and poll until completion."""
        response = self.kling_multi_image2video(**options)

        if not response.task_id:
            raise RuntimeError("No task ID returned from Kling multi-image2video request")

        return self._poll(
            lambda tid: self.kling_query_task("multi-image2video", tid),
            response.task_id,
            poll_interval,
            timeout,
        )

    def _poll(self, query_fn, task_id: str, poll_interval: float, timeout: float) -> TaskStatus:
        start = time.time()
        while True:
            status = query_fn(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)
