from typing import Literal

from bria_client import BriaAsyncClient
from mcp.server.fastmcp import Context
from mcp.types import CallToolResult, TextContent


async def remove_video_background(ctx: Context, video: str, background_color: Literal["Black", "White"] = "Black") -> CallToolResult:
    """
    Remove background from a video while maintaining the foreground subject.

    Args:
        video: Video URL (preferred), file path, or base64 data. URLs provide best performance.
        background_color: Color to replace the removed background with. Defaults to "black".

    Returns:
        CallToolResult containing the video URL with removed background.
    """
    client = BriaAsyncClient()
    endpoint = "/video/edit/remove_background"

    payload = {"video": video, "background_color": background_color}

    status_response = await client.submit(endpoint=endpoint, payload=payload)
    status_response.raise_for_status()
    response = await client.poll(status_response, timeout=60 * 5, interval=3)
    response.raise_for_status()

    assert response.result is not None

    return CallToolResult(
        content=[
            TextContent(type="text", text=f"Video with background removed: {response.result.video_url}"),
        ]
    )


async def increase_video_resolution(ctx: Context, video: str, desired_increase: Literal[2, 4] = 2) -> CallToolResult:
    """
    Increase the resolution of a video by a specified factor.

    Args:
        video: Video URL (preferred), file path, or base64 data. URLs provide best performance.
        desired_increase: Factor by which to increase the video resolution. Defaults to 2.

    Returns:
        CallToolResult containing the video URL with increased resolution.
    """
    client = BriaAsyncClient()
    endpoint = "/video/edit/increase_resolution"

    payload = {"video": video, "desired_increase": desired_increase}

    status_response = await client.submit(endpoint=endpoint, payload=payload)
    status_response.raise_for_status()
    response = await client.poll(status_response, timeout=60 * 5, interval=3)
    response.raise_for_status()

    assert response.result is not None

    return CallToolResult(
        content=[
            TextContent(type="text", text=f"Enhanced Video: {response.result.video_url}"),
        ]
    )


async def video_mask_by_prompt(ctx: Context, video: str, prompt: str) -> CallToolResult:
    """
    Mask a video by a prompt (will trim video to match proper length).

    Args:
        video: Video URL (preferred), file path, or base64 data. URLs provide best performance.
        prompt: Text prompt to guide the masking process.

    Returns:
        CallToolResult containing the video URL with masked content.
    """
    client = BriaAsyncClient()
    endpoint = "/video/segment/mask_by_prompt"

    payload = {"video": video, "prompt": prompt, "auto_trim": True}

    status_response = await client.submit(endpoint=endpoint, payload=payload)
    status_response.raise_for_status()
    response = await client.poll(status_response, timeout=60 * 5, interval=3)
    response.raise_for_status()

    assert response.result is not None

    return CallToolResult(
        content=[
            TextContent(type="text", text=f"Video mask: {response.result.video_url}"),
        ]
    )


async def erase_from_video(ctx: Context, video: str, mask_video: str, prompt: str | None = None) -> CallToolResult:
    """
    Erase an object from a video based on a mask video or given prompt.

    Args:
        video: Video URL (preferred), file path, or base64 data. URLs provide best performance.
        mask_video: Video URL (preferred), file path, or base64 data. URLs provide best performance.
        prompt: Text prompt to guide the masking process.

    Returns:
        CallToolResult containing the video URL with masked content.
    """
    client = BriaAsyncClient()
    if (mask_video is None) and (prompt is not None):
        client = BriaAsyncClient()
        endpoint = "/video/segment/mask_by_prompt"
        payload = {"video": video, "prompt": prompt, "auto_trim": True}
        status_response = await client.submit(endpoint=endpoint, payload=payload)
        status_response.raise_for_status()
        mask_response = await client.poll(status_response, timeout=60 * 5, interval=3)
        mask_response.raise_for_status()
        mask_video = mask_response.result.mask_url

    endpoint = "/video/edit/erase"
    payload = {"video": video, "mask": mask_video, "auto_trim": True}

    status_response = await client.submit(endpoint=endpoint, payload=payload)
    status_response.raise_for_status()
    response = await client.poll(status_response, timeout=60 * 5, interval=3)
    response.raise_for_status()

    assert response.result is not None

    return CallToolResult(
        content=[
            TextContent(type="text", text=f"Segmented Video: {response.result.video_url}"),
        ]
    )
