"""
Media processing utilities for the VT.ai application.
"""

import asyncio
import os
import tempfile
from io import BytesIO

import chainlit as cl
import litellm
from openai import OpenAI

# Update imports to use vtai namespace
from vtai.utils import llm_settings_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(
                            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(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 ""


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(
                    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:
        # Using wait_for to enforce a timeout
        vresponse = await asyncio.wait_for(
            litellm.acompletion(
                user=get_user_session_id(),
                model=vision_model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {"type": "image_url", "image_url": {"url": input_image}},
                        ],
                    }
                ],
                timeout=45.0,  # Add a specific timeout in litellm
            ),
            timeout=60.0,  # Overall operation timeout
        )

        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(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(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()
