"""
title: LibreTranslate Filter
description: Translates messages between users and assistants in a chat system using the LibreTranslate API.
author: @iamg30
author_url: https://openwebui.com/u/iamg30/
funding_url: https://github.com/open-webui
version: 0.2.1
license: MIT
"""

from typing import List, Optional, Callable, Awaitable
from pydantic import BaseModel, Field
import requests
import os

from utils.misc import get_last_user_message, get_last_assistant_message


class Filter:
    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for the filter operations."
        )
        libretranslate_url: str = Field(
            default="http://localhost:5000",
            description="URL for the LibreTranslate API",
        )
        source_user: str = Field(
            default="auto", description="Default source language for user messages"
        )
        target_user: str = Field(
            default="en", description="Default target language for user messages"
        )
        source_assistant: str = Field(
            default="en", description="Default source language for assistant messages"
        )
        target_assistant: str = Field(
            default="es", description="Default target language for assistant messages"
        )
        DISPLAY_EVENT_EMITTERS: bool = Field(
            default=False,
            description="Whether to display event emitters during translation.",
        )

    class UserValves(BaseModel):
        source_user: str = Field(
            default="auto",
            description="User-specific source language for user messages",
        )
        target_user: str = Field(
            default="en", description="User-specific target language for user messages"
        )
        source_assistant: str = Field(
            default="en",
            description="User-specific source language for assistant messages",
        )
        target_assistant: str = Field(
            default="es",
            description="User-specific target language for assistant messages",
        )
        DISPLAY_EVENT_EMITTERS: bool = Field(
            default=False,
            description="Whether to display event emitters during translation (user-specific setting).",
        )

    def __init__(self):
        self.file_handler = False
        self.valves = self.Valves()

    async def translate(
        self,
        text: str,
        source: str,
        target: str,
        __event_emitter__: Callable[[dict], Awaitable[None]],
        display_event_emitters: bool,
    ) -> str:
        payload = {
            "q": text,
            "source": source,
            "target": target,
        }

        if display_event_emitters:
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "status": "in_progress",
                        "description": f"Translating text from {source} to {target}",
                        "done": False,
                    },
                }
            )

        try:
            r = requests.post(
                f"{self.valves.libretranslate_url}/translate", json=payload
            )
            r.raise_for_status()

            data = r.json()
            translated_text = data["translatedText"]

            if display_event_emitters:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "status": "complete",
                            "description": f"Translation completed",
                            "done": True,
                        },
                    }
                )

            return translated_text
        except Exception as e:
            print(f"Error translating text: {e}")
            if display_event_emitters:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "status": "error",
                            "description": f"Error translating text: {e}",
                            "done": True,
                        },
                    }
                )
            return text

    async def inlet(
        self,
        body: dict,
        __event_emitter__: Callable[[dict], Awaitable[None]],
        __user__: dict = {},
    ) -> dict:
        print(f"inlet:{__name__}")

        # Get user_valves directly from __user__ dict
        user_valves = __user__.get("valves")
        if not user_valves:
            user_valves = self.UserValves()  # Use default UserValves if not provided

        # Determine whether to display event emitters based on both global and user-specific settings
        display_event_emitters = (
            self.valves.DISPLAY_EVENT_EMITTERS and user_valves.DISPLAY_EVENT_EMITTERS
        )

        messages = body["messages"]
        user_message = get_last_user_message(messages)

        print(f"User message: {user_message}")

        source_user = user_valves.source_user or self.valves.source_user
        target_user = user_valves.target_user or self.valves.target_user

        translated_user_message = await self.translate(
            user_message,
            source_user,
            target_user,
            __event_emitter__,
            display_event_emitters,
        )

        print(f"Translated user message: {translated_user_message}")

        for message in reversed(messages):
            if message["role"] == "user":
                message["content"] = translated_user_message
                break

        body = {**body, "messages": messages}
        return body

    async def outlet(
        self,
        body: dict,
        __event_emitter__: Callable[[dict], Awaitable[None]],
        __user__: dict = {},
    ) -> dict:
        print(f"outlet:{__name__}")

        # Get user_valves directly from __user__ dict
        user_valves = __user__.get("valves")
        if not user_valves:
            user_valves = self.UserValves()  # Use default UserValves if not provided

        # Determine whether to display event emitters based on both global and user-specific settings
        display_event_emitters = (
            self.valves.DISPLAY_EVENT_EMITTERS and user_valves.DISPLAY_EVENT_EMITTERS
        )

        messages = body["messages"]
        assistant_message = get_last_assistant_message(messages)

        print(f"Assistant message: {assistant_message}")

        source_assistant = user_valves.source_assistant or self.valves.source_assistant
        target_assistant = user_valves.target_assistant or self.valves.target_assistant

        translated_assistant_message = await self.translate(
            assistant_message,
            source_assistant,
            target_assistant,
            __event_emitter__,
            display_event_emitters,
        )

        print(f"Translated assistant message: {translated_assistant_message}")

        for message in reversed(messages):
            if message["role"] == "assistant":
                message["content"] = translated_assistant_message
                break

        body = {**body, "messages": messages}
        return body