"""
title: Better R1
author: Aaron Feng
author_url: https://github.com/AaronFeng753
version: 0.2
"""

from pydantic import BaseModel, Field
from typing import Optional
import re
import time


class Filter:
    start_time = None  # class variable to track time

    class Valves(BaseModel):
        priority: int = Field(default=0, description="")
        max_turns: int = Field(default=9999, description="")

    class UserValves(BaseModel):
        max_turns: int = Field(default=9999, description="")

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

    def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        messages = body.get("messages", [])
        modified_messages = []

        for message in messages:
            if message.get("role") == "assistant":
                message_content = message.get("content", "")
                pattern = r"<details[^>]*>.*?</details>"
                modified_content = re.sub(pattern, "", message_content, flags=re.DOTALL)
                modified_message = {"role": "assistant", "content": modified_content}
            else:
                modified_message = message
            modified_messages.append(modified_message)

        body["messages"] = modified_messages

        Filter.start_time = time.time()  # set initial time when instance created

        return body

    async def outlet(
        self, body: dict, __user__: Optional[dict] = None, __event_emitter__=None
    ) -> dict:
        elapsed_time = time.time() - Filter.start_time

        messages = body.get("messages", [])
        if not messages:
            return body

        model_response = messages[-1].get("content", "")

        model_response = model_response.replace(
            "<think>",
            f"<details>\n<summary>Thought for {int(elapsed_time)} seconds 💭</summary>",
        )
        model_response = model_response.replace("</think>", "\n\n---\n\n</details>")

        messages[-1]["content"] = model_response
        body["messages"] = messages
        
        Filter.start_time = time.time()  # set initial time when instance created

        return body
