"""
title: Add or delete text for every message
author: anfi
author_url: None
funding_url: None
version: 0.1
"""
# This can add a string of YOUR CUSTOM TEXT before and/or after your messages.
# (Write your strings of text in self.INLET_TEXT_PREPEND and self.INLET_TEXT_APPEND)
# Additionally, this removes anything once per returned message using a regular expression.
# In this example, it adds to the end of every message:
# an instruction for the model to think in a "thingking tag", which may theoretically
# improve performance, but it will definitely make it more expensive (thinking text eats tokens) and slower.
# It then on outlet - from model to you - removes anything in a thinking tag, so you see only clear messages.
# Alternatively, the same can be done to the AI's messages to you.

from typing import Optional, Dict, List
import json
from datetime import datetime
import os
import re
from pydantic import BaseModel, Field
class Filter:
    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for the filter operations."
        )
        pass
    
 
    def __init__(self) -> None:
        # Write your parameters here!: 
        self.DEBUG_FILE_PATH: str = "/app/backend/data/debuginfo.txt"
        self.ERROR_FILE_PATH: str = "/app/backend/data/debuginfo_errors.txt"
        self.DEBUG: bool = False
        self.INLET_TEXT_PREPEND: str = ""
        self.INLET_TEXT_APPEND: str = "Before answering, write thinking in a  tag, write better answer right away"
        self.INLET_TEXT_CUT_REGEXP: str = r"^$" # This regexp deletes nothing.
        self.OUTLET_TEXT_PREPEND: str = ""
        self.OUTLET_TEXT_APPEND: str = ""
        self.OUTLET_TEXT_CUT_REGEXP: str = r"[\s\S]*?" # This deletes thinking from model's outputs. It will only do it once per message

    def _prepend_text(self, messages: List[Dict[str, str]], text: str) -> None:
        if messages:
            messages[-1]["content"] = text + messages[-1]["content"]

    def _append_text(self, messages: List[Dict[str, str]], text: str) -> None:
        if messages:
            messages[-1]["content"] += text

    def _cut_text(self, messages: List[Dict[str, str]], regexp: str) -> None:
        if messages:
            messages[-1]["content"] = re.sub(regexp, "", messages[-1]["content"], count=1)

    def inlet(self, body: Dict[str, any], __user__: Optional[Dict[str, any]] = None) -> Dict[str, any]:
        # Process incoming messages and apply filtering.
        # The __user__ parameter is unused but kept for future modifications.

        # Args:
        #     body (Dict[str, any]): The incoming request body containing messages.
        #     __user__ (Optional[Dict[str, any]]): Unused parameter for potential future use.

        # Returns:
        #     Dict[str, any]: The processed request body with filtered messages.
        try:
            # Extract the original messages from the body
            original_messages: List[Dict[str, str]] = body.get("messages", [])

            # Apply text manipulations
            self._prepend_text(original_messages, self.INLET_TEXT_PREPEND)
            self._append_text(original_messages, self.INLET_TEXT_APPEND)
            self._cut_text(original_messages, self.INLET_TEXT_CUT_REGEXP)

            # Update the body with the manipulated messages
            body["messages"] = original_messages

            # If debug mode is enabled, log debug information
            if self.DEBUG:
                self._log_debug_info(
                    "inlet", body, __user__, self.DEBUG_FILE_PATH
                )

            # Return the updated body
            return body
        except Exception as e:
            # If an exception occurs and debug mode is enabled, log the error
            if self.DEBUG:
                self._log_debug_info(
                    "error", {"error": str(e)}, None, self.ERROR_FILE_PATH
                )
            # Return the original body in case of an error
            return body

    def outlet(self, body: Dict[str, any], __user__: Optional[Dict[str, any]] = None) -> Dict[str, any]:
        # Process outgoing messages and apply filtering.
        # The __user__ parameter is unused but kept for future modifications.

        # Args:
        #     body (Dict[str, any]): The outgoing request body containing messages.
        #     __user__ (Optional[Dict[str, any]]): Unused parameter for potential future use.

        # Returns:
        #     Dict[str, any]: The processed request body with filtered messages.
        try:
            # Extract the original messages from the body
            original_messages: List[Dict[str, str]] = body.get("messages", [])

            # Apply text manipulations
            self._prepend_text(original_messages, self.OUTLET_TEXT_PREPEND)
            self._append_text(original_messages, self.OUTLET_TEXT_APPEND)
            self._cut_text(original_messages, self.OUTLET_TEXT_CUT_REGEXP)

            # Update the body with the manipulated messages
            body["messages"] = original_messages

            # If debug mode is enabled, log debug information
            if self.DEBUG:
                self._log_debug_info(
                    "outlet", body, __user__, self.DEBUG_FILE_PATH
                )

            # Return the updated body
            return body
        except Exception as e:
            # If an exception occurs and debug mode is enabled, log the error
            if self.DEBUG:
                self._log_debug_info(
                    "error", {"error": str(e)}, None, self.ERROR_FILE_PATH
                )
            # Return the original body in case of an error
            return body

    def _log_debug_info(
        self,
        logging_source_method_name: str,
        body_content: Dict[str, any],
        user_info: Optional[Dict[str, any]],
        log_file_path: str
    ) -> None:
        try:
            os.makedirs(os.path.dirname(log_file_path), exist_ok=True)

            # Create a JSON structure for the log information
            log_data = {
                "timestamp": datetime.now().isoformat(),
                "method": logging_source_method_name,
                "body": body_content,  # Filtered body content
                "user": user_info
            }

            with open(log_file_path, "a") as log_file:
                # Write separator before the log entry
                log_file.write(f"\n{'='*50}\n")
                
                # Write the JSON structure
                json.dump(log_data, log_file, indent=2)
                
                # Write separator after the log entry
                log_file.write(f"\n{'='*50}\n")
        except Exception as e:
            print(f"Error writing debug info: {str(e)}")