"""
title: Google Translate Filter
author: justinh-rahb and OriginalSimon
author_url: https://github.com/justinh-rahb
funding_url: https://github.com/open-webui
version: 0.1.5
license: MIT
"""

import re
from typing import List, Optional
from pydantic import BaseModel
import requests
import logging

from utils.misc import get_last_user_message, get_last_assistant_message

# Configure logging
logging.basicConfig(level=logging.WARNING)  # Set the desired logging level


class Filter:
    class Valves(BaseModel):
        source_user: str = "auto"
        target_user: str = "en"
        source_assistant: str = "en"
        target_assistant: str = "en"

    def __init__(self) -> None:
        self.valves = self.Valves()
        self.code_blocks = []  # List to store code blocks

    def translate(self, text: str, source: str, target: str) -> str:
        url = "https://translate.googleapis.com/translate_a/single"
        params = {
            "client": "gtx",
            "sl": source,
            "tl": target,
            "dt": "t",
            "q": text,
        }

        try:
            r = requests.get(
                url, params=params, timeout=10
            )  # Add timeout for robustness
            r.raise_for_status()
            result = r.json()
            translated_text = "".join([sentence[0] for sentence in result[0]])
            return translated_text
        except requests.exceptions.RequestException as e:
            error_msg = f"Translation API error: {str(e)}"
            logging.error(error_msg)
            return f"{text}\n\n[Translation failed: {error_msg}]"
        except Exception as e:
            error_msg = f"Unexpected error during translation: {str(e)}"
            logging.exception(error_msg)  # Log traceback for unexpected errors
            return f"{text}\n\n[Translation failed: {error_msg}]"

    def split_text_around_table(self, text: str) -> List[str]:
        table_regex = r"((?:^.*?\|.*?\n)+)(?=\n[^\|\s].*?\|)"
        matches = re.split(table_regex, text, flags=re.MULTILINE)

        if len(matches) > 1:
            return [matches[0], matches[1]]
        else:
            return [text, ""]

    def clean_table_delimiters(self, text: str) -> str:
        # Replace multiple spaces around table delimiters with a single dash
        return re.sub(r"(\|\s*-+\s*)+", lambda m: m.group(0).replace(" ", "-"), text)

    async def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        print(f"inlet:{__name__}")
        print(f"source_user: {self.valves.source_user}")
        print(f"target_user: {self.valves.target_user}")
        print(f"source_assistant: {self.valves.source_assistant}")
        print(f"target_assistant: {self.valves.target_assistant}")

        user_message = get_last_user_message(body["messages"])

        # Find and store code blocks
        code_regex = r"```(.*?)```"
        self.code_blocks = re.findall(code_regex, user_message, flags=re.DOTALL)

        # Replace code blocks with placeholders temporarily
        user_message_processed = re.sub(
            code_regex, "__CODE_BLOCK__", user_message, flags=re.DOTALL
        )

        if self.valves.source_user != self.valves.target_user:
            parts = self.split_text_around_table(user_message_processed)
            text_before_table, table_text = parts

            translated_before_table = self.translate(
                text_before_table,
                self.valves.source_user,
                self.valves.target_user,
            )

            translated_user_message = translated_before_table + table_text
            translated_user_message = self.clean_table_delimiters(
                translated_user_message
            )

            # Restore code blocks in translated message
            for code in self.code_blocks:
                translated_user_message = translated_user_message.replace(
                    "__CODE_BLOCK__", f"```{code}```", 1
                )

            for message in reversed(body["messages"]):
                if message["role"] == "user":
                    if "[Translation failed:" in translated_user_message:
                        print(
                            f"Translation failed for language pair {self.valves.source_user} to {self.valves.target_user}"
                        )
                        # Optionally, you could decide not to update the message content if translation failed
                        # return body
                    else:
                        message["content"] = translated_user_message
                    break

        return body

    async def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        print(f"outlet:{__name__}")
        print(f"source_user: {self.valves.source_user}")
        print(f"target_user: {self.valves.target_user}")
        print(f"source_assistant: {self.valves.source_assistant}")
        print(f"target_assistant: {self.valves.target_assistant}")

        assistant_message = get_last_assistant_message(body["messages"])

        # Find and store code blocks
        code_regex = r"```(.*?)```"
        self.code_blocks = re.findall(code_regex, assistant_message, flags=re.DOTALL)

        # Replace code blocks with placeholders temporarily
        assistant_message_processed = re.sub(
            code_regex, "__CODE_BLOCK__", assistant_message, flags=re.DOTALL
        )

        if self.valves.source_assistant != self.valves.target_assistant:
            parts = self.split_text_around_table(assistant_message_processed)
            text_before_table, table_text = parts

            translated_before_table = self.translate(
                text_before_table,
                self.valves.source_assistant,
                self.valves.target_assistant,
            )

            translated_assistant_message = translated_before_table + table_text
            translated_assistant_message = self.clean_table_delimiters(
                translated_assistant_message
            )

            # Restore code blocks in translated message
            for code in self.code_blocks:
                translated_assistant_message = translated_assistant_message.replace(
                    "__CODE_BLOCK__", f"```{code}```", 1
                )

            for message in reversed(body["messages"]):
                if message["role"] == "assistant":
                    if "[Translation failed:" in translated_assistant_message:
                        print(
                            f"Translation failed for language pair {self.valves.source_assistant} to {self.valves.target_assistant}"
                        )
                        # Optionally, you could decide not to update the message content if translation failed
                        # return body
                    else:
                        message["content"] = translated_assistant_message
                    break

        return body