"""VideoProvider implementation for ByteDance SeeDance video generation API."""

__author__ = "Lene Preuss <lene.preuss@gmail.com>"

import base64
import logging
import mimetypes
import os
from datetime import datetime, timezone
from pathlib import Path

from byteplussdkarkruntime import Ark  # type: ignore

from stable_delusion.exceptions import AuthenticationError, VideoGenerationError
from stable_delusion.providers.video_provider import VideoProvider
from stable_delusion.providers.types import (
    Cost,
    TaskHandle,
    TaskStatus,
    TaskState,
    VideoGenParams,
    Capabilities,
)


_COST_TABLE = {
    "seedance-1-5-pro-251215": 0.24,
    "seedance-1-0-pro-fast-251015": 0.07,
}
_DEFAULT_COST = 0.24

_STATUS_MAP = {
    "queued": TaskState.QUEUED,
    "running": TaskState.RUNNING,
    "succeeded": TaskState.SUCCEEDED,
    "failed": TaskState.FAILED,
    "cancelled": TaskState.CANCELLED,
}


def _image_to_data_uri(image_path: Path) -> str:
    """Encode an image file as a base64 data URI with correct MIME type."""
    mime_type, _ = mimetypes.guess_type(str(image_path))
    if mime_type is None:
        mime_type = "image/jpeg"
    with open(image_path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")
    return f"data:{mime_type};base64,{encoded}"


class SeedanceProvider(VideoProvider):
    """VideoProvider for ByteDance SeeDance video generation API."""

    def __init__(self, api_key: str) -> None:
        self._client = Ark(api_key=api_key)

    @classmethod
    def create_with_env_key(cls, env_var: str = "ARK_API_KEY") -> "SeedanceProvider":
        """Create provider using API key from environment variable."""
        api_key = os.getenv(env_var)
        if not api_key:
            raise AuthenticationError(
                f"BytePlus ARK API key not found in environment variable: {env_var}"
            )
        return cls(api_key)

    @property
    def name(self) -> str:
        return "seedance"

    @property
    def capabilities(self) -> Capabilities:
        return Capabilities(
            text_to_media=True,
            media_to_media=True,
            max_duration=12,
            supported_resolutions=("720p",),
            supported_ratios=("16:9", "9:16", "1:1"),
            supports_audio=True,
        )

    def submit(self, prompt: str, keyframes: list[Path], params: VideoGenParams) -> TaskHandle:
        """Submit a video generation task and return a handle for polling."""
        logging.info("Submitting SeeDance video generation task")
        logging.debug(
            "Prompt: %s, keyframes: %d, duration: %ds", prompt, len(keyframes), params.duration
        )

        request_params = self._build_request_params(prompt, keyframes, params)

        try:
            create = self._client.content_generation.tasks.create  # pylint: disable=no-member
            # pylint: disable=unexpected-keyword-arg
            task = create(**request_params)  # type: ignore[call-arg]
            # pylint: enable=unexpected-keyword-arg
        except Exception as e:
            raise VideoGenerationError(f"SeeDance task submission failed: {e}") from e

        logging.info("Task submitted: %s", task.id)
        return TaskHandle(
            task_id=task.id,
            provider="seedance",
            submitted_at=datetime.now(timezone.utc),
            prompt=prompt,
            params=params,
        )

    def poll(self, handle: TaskHandle) -> TaskStatus:
        """Poll the status of a submitted task."""
        logging.debug("Polling task: %s", handle.task_id)

        try:
            response = self._client.content_generation.tasks.get(task_id=handle.task_id)
        except Exception as e:
            raise VideoGenerationError(
                f"Failed to poll task {handle.task_id}: {e}",
                task_id=handle.task_id,
            ) from e

        state = _STATUS_MAP.get(response.status, TaskState.FAILED)

        if state == TaskState.RUNNING:
            return self._parse_running_status(response)
        if state == TaskState.SUCCEEDED:
            return self._parse_succeeded_status(response, handle)
        if state == TaskState.FAILED:
            return self._parse_failed_status(response)

        return TaskStatus(state=state)

    def estimate_cost(self, params: VideoGenParams) -> Cost:
        """Return estimated cost for the given video generation params."""
        amount = _COST_TABLE.get(params.model, _DEFAULT_COST)
        return Cost(
            amount=amount,
            currency="USD",
            input_tokens=0,
            output_tokens=0,
            model=params.model,
            timestamp=datetime.now(timezone.utc),
        )

    def _build_request_params(
        self, prompt: str, keyframes: list[Path], params: VideoGenParams
    ) -> dict:
        request_params: dict = {
            "model": params.model,
            "content": [{"type": "text", "text": prompt}],
            "extra": {
                "resolution": params.resolution,
                "ratio": params.ratio,
                "duration": str(params.duration),
                "generate_audio": params.generate_audio,
                "return_last_frame": params.return_last_frame,
            },
        }

        for i, kf in enumerate(keyframes[:2]):
            key = "first_frame" if i == 0 else "last_frame"
            request_params["content"].append(
                {
                    "type": "image_url",
                    "image_url": {"url": _image_to_data_uri(kf)},
                    "role": key,
                }
            )

        return request_params

    def _parse_running_status(self, response: object) -> TaskStatus:
        progress_raw = getattr(response, "progress", None)
        progress = progress_raw / 100.0 if progress_raw is not None else None
        eta = getattr(response, "estimated_time", None)
        return TaskStatus(state=TaskState.RUNNING, progress=progress, eta_seconds=eta)

    def _parse_succeeded_status(self, response: object, handle: TaskHandle) -> TaskStatus:
        content = getattr(response, "content", None)
        video_url = None
        last_frame_url = None
        if content:
            video = getattr(content, "video", None)
            if video:
                video_url = getattr(video, "url", None)
            last_frame_url = getattr(content, "last_frame_url", None)

        usage = getattr(response, "usage", None)
        cost = None
        if usage:
            input_tokens = getattr(usage, "input_tokens", 0) or 0
            output_tokens = getattr(usage, "output_tokens", 0) or 0
            amount = _COST_TABLE.get(handle.params.model, _DEFAULT_COST)
            cost = Cost(
                amount=amount,
                currency="USD",
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                model=handle.params.model,
                timestamp=datetime.now(timezone.utc),
            )

        return TaskStatus(
            state=TaskState.SUCCEEDED,
            video_url=video_url,
            last_frame_url=last_frame_url,
            cost=cost,
        )

    def _parse_failed_status(self, response: object) -> TaskStatus:
        error_msg = None
        error_obj = getattr(response, "error", None)
        if error_obj:
            error_msg = getattr(error_obj, "message", str(error_obj))
        return TaskStatus(state=TaskState.FAILED, error=error_msg)
