from __future__ import annotations

import asyncio
from typing import Any, Dict, List

from core.tool import ToolBase
from utils.common_util import extract_json
from utils.openai_util import OpenAITool
from core.config import LLM, MAX_WORKERS


class LLMTool(ToolBase):
    """LLM-backed processor: builds prompts, calls LLM, parses JSON."""

    def __init__(
        self, system_prompt: str, llm_key: str = "fastgpt_segment_review"
    ) -> None:
        super().__init__()
        self.system_prompt = system_prompt
        self.llm = OpenAITool(LLM[llm_key], max_workers=MAX_WORKERS)

    def build_messages(self, user_content: str, system_content: str = None) -> List[Dict[str, str]]:
        return [
            {"role": "system", "content": system_content or self.system_prompt},
            {"role": "user", "content": user_content},
        ]

    async def chat_async(self, messages: List[Dict[str, str]]):
        return await self.llm.chat(messages)

    async def chat_batch_async(self, messages_list: List[List[Dict[str, str]]]):
        return await self.llm.mul_chat(messages_list)

    def run_with_loop(self, coro):
        try:
            return asyncio.run(coro)
        except RuntimeError as e:
            print(
                f"RuntimeError in run_with_loop: {e}, trying to get event loop and run until complete."
            )
            loop = asyncio.get_event_loop()
            return loop.run_until_complete(coro)

    def parse_first_json(self, resp: str) -> Dict[str, Any]:
        try:
            data = extract_json(resp)
            return data[0] if data else {}
        except Exception:
            return {}
