from flowllm.core.context import C
from flowllm.core.enumeration import Role
from flowllm.core.op import BaseAsyncToolOp
from flowllm.core.schema import ToolCall, Message

from finmcp.core.utils import get_datetime


@C.register_op()
class ExtractLongTextOp(BaseAsyncToolOp):
    file_path: str = __file__

    def __init__(self, max_content_char_length: int = 50000, **kwargs):
        super().__init__(**kwargs)
        self.max_content_char_length = max_content_char_length

    def build_tool_call(self) -> ToolCall:
        return ToolCall(
            **{
                "description": "Utilize the capabilities of LLM to parse content relevant to the query from long_text.",
                "input_schema": {
                    "long_text": {
                        "type": "string",
                        "description": "long_text",
                        "required": True,
                    },
                    "query": {
                        "type": "string",
                        "description": "query",
                        "required": True,
                    },
                },
            },
        )

    async def async_execute(self):
        long_text: str = self.input_dict["long_text"]
        long_text = long_text[: self.max_content_char_length]
        query: str = self.input_dict["query"]

        extract_content_prompt = self.prompt_format(
            prompt_name="extract_content_prompt",
            long_text=long_text,
            datetime=get_datetime(),
            query=query,
        )
        assistant_message = await self.llm.achat(messages=[Message(role=Role.USER, content=extract_content_prompt)])
        self.set_output(assistant_message.content)
