"""
Media processing utilities for the VT application.

Handles image, audio, and text-to-speech processing.
"""

import asyncio
import base64
import os
import tempfile
from io import BytesIO
from pathlib import Path

import chainlit as cl
import litellm
from openai import OpenAI
from PIL import Image

# Update imports to use vtai namespace
from vtai.utils import llm_providers_config as conf
from vtai.utils.config import logger
from vtai.utils.user_session_helper import (
    get_setting,
    get_user_session_id,
    update_message_history_from_assistant,
)


async def handle_tts_response(context: str, openai_client: OpenAI) -> None:
    """
    Generates and sends a TTS audio response using OpenAI's Audio API.

    Args:
        context: Text to convert to speech
        openai_client: OpenAI client instance
    """
    enable_tts_response = get_setting(conf.SETTINGS_ENABLE_TTS_RESPONSE)
    if enable_tts_response is False or not context:
        return

    model = get_setting(conf.SETTINGS_TTS_MODEL)
    voice = get_setting(conf.SETTINGS_TTS_VOICE_PRESET_MODEL)

    try:
        temp_filepath = os.path.join(tempfile.gettempdir(), "tts-output.mp3")

        # Using a custom timeout for the TTS request to avoid hanging connections
        with openai_client.audio.speech.with_streaming_response.create(
            model=model, voice=voice, input=context
        ) as response:
            response.stream_to_file(temp_filepath)

            # Allow a small delay for file operations to complete
            await asyncio.sleep(0.1)

            if os.path.exists(temp_filepath):
                await cl.Message(
                    author=model,
                    content="",
                    elements=[
                        cl.Audio(path=temp_filepath, display="inline"),
                        cl.Text(
                            name="TTS Info",
                            content=f"You're hearing an AI voice generated by OpenAI's {model} model, using the {voice} style. You can customize this in Settings if you'd like!",
                            display="inline",
                        ),
                    ],
                ).send()

                update_message_history_from_assistant(context)
            else:
                logger.warning("TTS file was not created successfully")

    except asyncio.CancelledError:
        logger.warning("TTS operation was cancelled")
        # Re-raise to ensure proper cleanup
        raise
    except Exception as e:
        logger.error(f"Error generating TTS response: {e}")
        await cl.Message(content=f"Failed to generate speech: {str(e)}").send()


async def handle_audio_transcribe(
    path: str, audio_file: BytesIO, openai_client: OpenAI
) -> str:
    """
    Transcribes audio to text using OpenAI's Whisper model.

    Args:
        path: Path to the audio file
        audio_file: BytesIO object for the audio file
        openai_client: OpenAI client instance

    Returns:
        The transcribed text
    """
    model = conf.DEFAULT_WHISPER_MODEL
    try:
        # Add a timeout to the transcription request
        transcription = await asyncio.wait_for(
            openai_client.audio.transcriptions.create(model=model, file=audio_file),
            timeout=30.0,  # 30 second timeout
        )
        text = transcription.text

        await cl.Message(
            content="",
            author=model,
            elements=[
                cl.Audio(path=path, display="inline"),
                cl.Text(name="Transcription", content=text, display="inline"),
            ],
        ).send()

        update_message_history_from_assistant(text)
        return text
    except asyncio.TimeoutError:
        logger.error("Audio transcription request timed out")
        await cl.Message(
            content="Audio transcription timed out. Please try with a shorter audio file."
        ).send()
        return ""
    except asyncio.CancelledError:
        logger.warning("Audio transcription was cancelled")
        raise
    except Exception as e:
        logger.error(f"Error transcribing audio: {e}")
        await cl.Message(content=f"Failed to transcribe audio: {str(e)}").send()
        return ""


def encode_image_to_base64(image_path):
    """
    Encodes an image file to base64 string.

    Args:
        image_path: Path to the image file

    Returns:
        Base64 encoded string with data URI prefix
    """
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode("utf-8")

            # Get image format
            img_format = Path(image_path).suffix.lstrip(".").lower()
            if img_format == "jpg":
                img_format = "jpeg"

            # Return with proper data URI format
            return f"data:image/{img_format};base64,{encoded_string}"
    except Exception as e:
        logger.error(f"Error encoding image to base64: {e}")
        raise e


async def handle_vision(
    input_image: str,
    prompt: str,
    is_local: bool = False,
) -> None:
    """
    Handles vision processing tasks using the specified vision model.
    Sends the processed image and description to the user.

    Args:
        input_image: Path or URL to the image
        prompt: Text prompt to accompany the image
        is_local: Whether the image is a local file or URL
    """
    vision_model = (
        conf.DEFAULT_VISION_MODEL
        if is_local
        else get_setting(conf.SETTINGS_VISION_MODEL)
    )

    supports_vision = litellm.supports_vision(model=vision_model)

    if supports_vision is False:
        logger.warning(f"Unsupported vision model: {vision_model}")
        await cl.Message(
            content="",
            elements=[
                cl.Text(
                    name="Vision Model Warning",
                    content=f"It seems the vision model `{vision_model}` doesn't support image processing. Please choose a different model in Settings that offers Vision capabilities.",
                    display="inline",
                )
            ],
        ).send()
        return

    message = cl.Message(
        content="I'm analyzing the image. This might take a moment.",
        author=vision_model,
    )

    await message.send()
    try:
        # For local images, convert to base64 if using Gemini
        image_content = input_image
        if is_local and "gemini" in vision_model.lower():
            logger.info(
                f"Converting local image to base64 for Gemini model: {vision_model}"
            )
            try:
                image_content = encode_image_to_base64(input_image)
            except Exception as e:
                logger.error(f"Failed to encode image to base64: {e}")
                await cl.Message(
                    content=f"Failed to process the image: {str(e)}"
                ).send()
                return

        # Prepare message content format for image
        input_messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": image_content}},
                ],
            }
        ]

        # Use the Chat Completions API directly instead of trying Response API first
        logger.info(f"Using Chat Completions API for vision model: {vision_model}")
        vresponse = await asyncio.wait_for(
            litellm.acompletion(
                user=get_user_session_id(),
                model=vision_model,
                messages=input_messages,
                timeout=45.0,  # Add a specific timeout in litellm
                response_format={"type": "text"},
            ),
            timeout=60.0,  # Overall operation timeout
        )

        # With the Chat Completions API format
        description = vresponse.choices[0].message.content

        if is_local:
            image = cl.Image(path=input_image, display="inline")
        else:
            image = cl.Image(url=input_image, display="inline")

        message = cl.Message(
            author=vision_model,
            content="",
            elements=[
                image,
                cl.Text(name=vision_model, content=description, display="inline"),
            ],
            actions=[
                cl.Action(
                    icon="speech",
                    name="speak_chat_response_action",
                    payload={"value": description},
                    label="Speak response",
                )
            ],
        )

        update_message_history_from_assistant(description)
        await message.send()
    except asyncio.TimeoutError:
        logger.error("Vision processing request timed out")
        await cl.Message(
            content="Image analysis timed out. The image might be too complex or the service is busy."
        ).send()
    except asyncio.CancelledError:
        logger.warning("Vision processing was cancelled")
        raise
    except Exception as e:
        logger.error(f"Error processing image with vision model: {e}")
        await cl.Message(content=f"Failed to analyze the image: {str(e)}").send()


async def handle_trigger_async_image_gen(query: str) -> None:
    """
    Triggers asynchronous image generation using the default image generation model.

    Args:
        query: Text prompt for image generation
    """
    image_gen_model = conf.DEFAULT_IMAGE_GEN_MODEL
    update_message_history_from_assistant(query)

    message = cl.Message(
        content="Sure! I'll create an image based on your description. This might take a moment, please be patient.",
        author=image_gen_model,
    )
    await message.send()

    style = get_setting(conf.SETTINGS_IMAGE_GEN_IMAGE_STYLE)
    quality = get_setting(conf.SETTINGS_IMAGE_GEN_IMAGE_QUALITY)
    try:
        # Using wait_for to enforce a timeout
        image_response = await asyncio.wait_for(
            litellm.aimage_generation(
                user=get_user_session_id(),
                prompt=query,
                model=image_gen_model,
                style=style,
                quality=quality,
                timeout=45.0,  # Add a specific timeout
            ),
            timeout=60.0,  # Overall operation timeout
        )

        image_gen_data = image_response["data"][0]
        image_url = image_gen_data["url"]
        revised_prompt = image_gen_data.get("revised_prompt", query)

        message = cl.Message(
            author=image_gen_model,
            content="Here's the image, along with a refined description based on your input:",
            elements=[
                cl.Image(url=image_url, display="inline"),
                cl.Text(
                    name=image_gen_model,
                    content=revised_prompt,
                    display="inline",
                ),
            ],
            actions=[
                cl.Action(
                    icon="speech",
                    name="speak_chat_response_action",
                    payload={"value": revised_prompt},
                    tooltip="Speak response",
                    label="Speak response",
                )
            ],
        )

        update_message_history_from_assistant(revised_prompt)
        await message.send()
    except asyncio.TimeoutError:
        logger.error("Image generation request timed out")
        await cl.Message(
            content="Image generation timed out. Please try a simpler description or try again later."
        ).send()
    except asyncio.CancelledError:
        logger.warning("Image generation was cancelled")
        raise
    except Exception as e:
        logger.error(f"Error generating image: {e}")
        await cl.Message(content=f"Failed to generate image: {str(e)}").send()
