# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description: WebResearcher Agent implementing the IterResearch paradigm.
"""
import json5
import re
import datetime
import asyncio
import random
import time

from typing import Any, Callable, Dict, List, Optional
from openai import OpenAI, APIError, APIConnectionError, APITimeoutError
import inspect
import sys
sys.path.append('..')
from webresearcher.base import Message, build_text_completion_prompt, count_tokens as count_tokens_base
from webresearcher.log import logger
from webresearcher.prompt import get_iterresearch_system_prompt
from webresearcher.tool_file import FileParser
from webresearcher.tool_scholar import Scholar
from webresearcher.tool_python import PythonInterpreter
from webresearcher.tool_search import Search
from webresearcher.tool_visit import Visit
from webresearcher.config import (
    LLM_API_KEY, 
    LLM_BASE_URL, 
    OBS_START, 
    OBS_END, 
    MAX_LLM_CALL_PER_RUN, 
    AGENT_TIMEOUT, 
    FILE_DIR,
    LLM_MODEL_NAME
)


TOOL_CLASS = [
    FileParser(),
    Scholar(),
    Visit(),
    Search(),
    PythonInterpreter(),
]
TOOL_MAP = {tool.name: tool for tool in TOOL_CLASS}


def today_date():
    return datetime.date.today().strftime("%Y-%m-%d")


class ResearchRound:
    """
    实现了 IterResearch 范式的核心状态管理器。
    
    状态只包含用于构建下一个提示的核心信息：
    - question (Q): 原始问题
    - current_report (R_{i-1}): 上一轮生成的报告
    - last_observation (O_{i-1}): 上一轮工具调用的结果
    """

    def __init__(self, question: str):
        self.question = question

        # R_{i-1}: 上一轮生成的报告，初始为空
        self.current_report = "This is the first round. The report is empty."

        # O_{i-1}: 上一轮工具调用的结果，初始为空
        self.last_observation = "This is the first round. No tool has been called yet."

    def get_context(self, system_prompt: str) -> List[Dict]:
        """
        构建当前轮次的精简上下文。
        这代表了论文中的 Workspace，即状态 s_t = (Q, R_{i-1}, O_{i-1})
        """
        user_content = (
            f"**Question:** {self.question}\n\n"
            f"**Current Report (R_{{i-1}}):**\n{self.current_report}\n\n"
            f"**Last Observation (O_{{i-1}}):**\n{self.last_observation}"
        )

        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content}
        ]


class WebResearcherAgent:
    """
    Web Research Agent implementing the IterResearch paradigm.
    
    The IterResearch paradigm uses a cyclical process of:
    1. Workspace (question + previous report + last tool result)
    2. Plan-Action (LLM generates plan and action)
    3. Tool Execution (execute the tool and get result)
    4. Synthesis (update report by integrating new information)
    
    This agent is fully independent and does not inherit from external frameworks.
    """

    def __init__(
            self,
            llm_config: Optional[Dict] = None,
            function_list: Optional[List[str]] = None,
            instruction: str = "",
            api_key: Optional[str] = None,
            base_url: Optional[str] = None,
            model: Optional[str] = None,
    ):
        llm_config = dict(llm_config or {})
        if api_key:
            llm_config["api_key"] = api_key
        if base_url:
            llm_config["base_url"] = base_url

        self.llm_config = llm_config
        self.llm_generate_cfg = self.llm_config.get("generate_cfg", {})
        self.model = model or self.llm_config.get("model", LLM_MODEL_NAME)  # 主模型
        self.api_key = self.llm_config.get("api_key", LLM_API_KEY)
        self.base_url = self.llm_config.get("base_url", LLM_BASE_URL)
        self.max_input_tokens = self.llm_config.get("max_input_tokens", 32000)
        self.llm_timeout = self.llm_config.get("llm_timeout", 300.0)
        self.agent_timeout = self.llm_config.get("agent_timeout", 600.0)
        self.function_list = function_list or list(TOOL_MAP.keys())
        self.instruction = instruction

    def parse_output(self, text: str) -> Dict[str, str]:
        """
        解析 LLM 的单次输出，严格提取 <plan>, <report>, 和 (<tool_call> 或 <answer> 或 <terminate>)。
        这是 IterResearch 范式的核心：LLM 在单次调用中生成三段式输出。
        """
        output = {
            "plan": "",
            "report": "",
            "tool_call": "",
            "answer": "",
            "terminate": False,
            "terminate_reason": "",
        }

        def _extract_last_block(pattern: str) -> str:
            matches = re.findall(pattern, text, flags=re.DOTALL | re.MULTILINE)
            for m in reversed(matches):
                if m and m.strip():
                    return m.strip()
            return ""

        # 1. 提取 <plan>
        plan = _extract_last_block(r"^\s*<plan>(.*?)</plan>")
        if plan:
            output["plan"] = plan

        # 2. 提取 <report>
        output["report"] = _extract_last_block(r"^\s*<report>(.*?)</report>")

        # 3. 提取 <tool_call>、<answer>、<terminate>
        output["tool_call"] = _extract_last_block(r"^\s*<tool_call>(.*?)</tool_call>")
        output["answer"] = _extract_last_block(r"^\s*<answer>(.*?)</answer>")
        term_body = _extract_last_block(r"^\s*<terminate>(.*?)</terminate>")
        if term_body != "":
            output["terminate"] = True
            output["terminate_reason"] = term_body

        if not output["tool_call"] and not output["answer"] and not output["terminate"]:
            logger.warning("LLM output did not contain <tool_call>, <answer>, or <terminate> tag.")

        return output

    async def call_server(self, msgs: List[Dict], stop_sequences: List[str] = None,
                          max_tries: int = 1) -> str:
        """异步方法，并使用 run_in_executor 处理同步的 OpenAI 库"""
        client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
            timeout=self.llm_timeout,
        )

        base_sleep_time = 1
        loop = asyncio.get_event_loop()

        stop_sequences = stop_sequences or ["<tool_response>"]

        for attempt in range(max_tries):
            try:
                request_params = {
                    "model": self.model,
                    "messages": msgs,
                    "stop": stop_sequences,
                    "temperature": self.llm_generate_cfg.get('temperature', 0.6),
                    "top_p": self.llm_generate_cfg.get('top_p', 0.95),
                    "presence_penalty": self.llm_generate_cfg.get('presence_penalty', 1.1)
                }
                model_thinking_type = self.llm_generate_cfg.get("model_thinking_type", "")
                if model_thinking_type:
                    request_params["extra_body"] = {
                        "thinking": {
                            "type": model_thinking_type,
                        }
                    }
                # [关键] 使用 run_in_executor 在线程池中运行同步的 blocking I/O
                chat_response = await loop.run_in_executor(
                    None,  # 使用默认的 ThreadPoolExecutor
                    lambda: client.chat.completions.create(**request_params)
                )

                content = chat_response.choices[0].message.content
                reasoning_content = None
                if hasattr(chat_response.choices[0].message, 'reasoning_content') and chat_response.choices[
                    0].message.reasoning_content:
                    reasoning_content = chat_response.choices[0].message.reasoning_content
                logger.debug(f"input messages: {msgs}, \nreasoning_content: {reasoning_content}, \nLLM Response: {content}")
                if content and content.strip():
                    return content.strip()
                else:
                    logger.warning(f"Attempt {attempt + 1}: Empty response received.")

            except (APIError, APIConnectionError, APITimeoutError) as e:
                logger.warning(
                    f"Attempt {attempt + 1} API error: {e}, base_url: {self.base_url}, api_key: {self.api_key}, model: {self.model}")
            except Exception as e:
                logger.error(f"Attempt {attempt + 1} unexpected error: {e}")

            if attempt < max_tries - 1:
                sleep_time = base_sleep_time * (2 ** attempt) + random.uniform(0, 1)
                sleep_time = min(sleep_time, 30)
                logger.warning(f"Retrying in {sleep_time:.2f}s...")
                await asyncio.sleep(sleep_time)  # [关键] 使用 await asyncio.sleep
            else:
                logger.error("All retry attempts exhausted. The LLM call failed.")
        return "LLM server error."

    def count_tokens(self, messages, model=None):
        if model is None:
            model = LLM_MODEL_NAME
        """Count tokens in messages"""
        try:
            # Convert dict messages to Message objects if needed
            full_message = []
            for x in messages:
                if isinstance(x, dict):
                    full_message.append(Message(**x))
                elif isinstance(x, Message):
                    full_message.append(x)
                else:
                    full_message.append(x)

            full_prompt = build_text_completion_prompt(full_message, allow_special=True)
            return count_tokens_base(full_prompt, model)
        except Exception as e:
            logger.warning(f"Failed to count tokens: {e}. Using simple split.")
            return sum(len(str(x).split()) for x in messages)

    async def custom_call_tool(self, tool_call_str: str) -> str:
        """async方法，并正确处理同步/异步工具"""
        loop = asyncio.get_event_loop()

        try:
            # 1. 优先检查是否包含 <code> 标签（处理 Python 代码）
            # 支持多种格式：
            # - python\n<code>...</code>
            # - <code>...</code>
            # - {"name": "python", ...}\n<code>...</code>
            if "<code>" in tool_call_str and "</code>" in tool_call_str:
                code_raw = tool_call_str.split("<code>", 1)[1].rsplit("</code>", 1)[0].strip()
                result = await loop.run_in_executor(None, TOOL_MAP['python'].call, code_raw)
                return result

            # 2. 处理 JSON 工具调用
            tool_call = json5.loads(tool_call_str)
            tool_name = tool_call.get('name', '')
            tool_args = tool_call.get('arguments', {})

            if tool_name not in TOOL_MAP:
                return f"Error: Tool {tool_name} not found"

            tool = TOOL_MAP[tool_name]

            # 3. [关键] 区分同步和异步工具
            if asyncio.iscoroutinefunction(tool.call):
                # 如果工具本身是 async (例如 FileParser)
                if tool_name == "parse_file":
                    params = {"files": tool_args.get("files")}
                    result = await tool.call(params, file_root_path=FILE_DIR)
                else:
                    result = await tool.call(tool_args)  # 假设其他 async 工具
            else:
                # 如果工具是 sync (例如 Search, Visit, Scholar)
                # 在 executor 中运行
                result = await loop.run_in_executor(None, tool.call, tool_args)

            return str(result) if not isinstance(result, str) else result

        except Exception as e:
            logger.error(f"Tool call parsing or execution failed: {e}")
            return f"Error: Tool call failed. Input: {tool_call_str}. Error: {e}"

    async def run(self, question, progress_callback: Optional[Callable[[Dict[str, Any]], Any]] = None):
        """
        严格按照 IterResearch 范式执行研究（单 LLM 调用）。
        
        循环如下:
        s_t = (Q, R_{i-1}, O_{i-1})
        LLM(s_t) -> (Plan_i, Report_i, Action_i)
        
        if Action_i == <answer>:
            STOP
        else:
            O_i = Tool(Action_i)
            s_{t+1} = (Q, R_i, O_i)
            LOOP
        """

        async def emit(event: Dict[str, Any]):
            if not callable(progress_callback):
                return
            event.setdefault("timestamp", datetime.datetime.utcnow().isoformat())
            try:
                maybe = progress_callback(event)
                if inspect.isawaitable(maybe):
                    await maybe
            except Exception as callback_err:
                logger.warning(f"progress_callback raised error: {callback_err}")

        start_time = time.time()

        # 1. 初始化研究轮次
        research_round = ResearchRound(question=question)
        system_prompt = get_iterresearch_system_prompt(today_date(), self.function_list, self.instruction, question=question)

        # 完整轨迹日志（用于调试）
        full_trajectory_log = []
        prediction = ''
        termination = ''

        num_llm_calls_available = MAX_LLM_CALL_PER_RUN
        round_num = 0

        while num_llm_calls_available > 0:
            if time.time() - start_time > self.agent_timeout:
                logger.warning("Agent timeout reached.")
                termination = "timeout"
                prediction = "No answer found (timeout)."
                break

            round_num += 1
            num_llm_calls_available -= 1

            # 2. 构建提示 (s_t = Q, R_{i-1}, O_{i-1})
            current_context = research_round.get_context(system_prompt)

            if round_num == 1:
                full_trajectory_log.extend(current_context)  # 仅记录初始上下文

            # 3. 单次 LLM 调用 (生成 P_i, R_i, A_i)
            content = ''
            request_msgs = ''
            is_last_call = False
            try:
                logger.debug(f"Round {round_num}: Calling LLM. Remaining calls: {num_llm_calls_available}")

                # If this is the last available LLM call, force final answer generation instead of further tool calls
                is_last_call = (num_llm_calls_available == 0)
                if is_last_call:
                    finalize_instruction = (
                        "You have reached the maximum allowed LLM calls for this run. "
                        "Do not call tools anymore. Based on your current report and the information gathered so far, "
                        "provide the final answer now in the three-part format: "
                        "<plan>...</plan> <report>...</report> <answer>...</answer>"
                    )
                    request_msgs = current_context + [{"role": "user", "content": finalize_instruction}]
                else:
                    request_msgs = current_context

                content = await self.call_server(request_msgs)

                full_trajectory_log.append({"role": "assistant", "content": content})
                logger.debug(f'Round {round_num} LLM response received.')
            except Exception as e:
                logger.error(f"Unknown Error: {e}")
                prediction = f"Error: Unknown {e}"
                termination = 'unknown error'
                break

            # 4. 解析 LLM 的结构化输出 (P_i, R_i, A_i / Answer_i)
            parsed = self.parse_output(content)
            plan_content = parsed["plan"]
            report_content = parsed["report"]
            action_content = parsed["tool_call"]
            answer_content = parsed["answer"]
            terminate_flag = parsed.get("terminate", False)
            terminate_reason = parsed.get("terminate_reason", "").strip()

            logger.debug(f"Round {round_num} - Plan: {plan_content}")
            logger.debug(f"Round {round_num} - Report: {report_content}")
            logger.debug(f"Round {round_num} - Action: {action_content}")
            logger.debug(f"Round {round_num} - Answer: {answer_content}")
            if terminate_flag:
                logger.debug(f"Round {round_num} - Terminate signaled. Reason: {terminate_reason}")

            await emit({
                "type": "round",
                "round": round_num,
                "plan": plan_content,
                "report": report_content,
                "action": action_content,
                "answer": answer_content,
                "terminate": terminate_flag,
            })

            # 5. 状态更新 (s_t -> s_{t+1})

            # 5.1 更新报告 (R_i)
            # 无论 LLM 接下来是调用工具还是回答，它都必须生成一份新报告。
            # 这份新报告 R_i 将用于 s_{t+1}
            if report_content:
                research_round.current_report = report_content
            else:
                # 如果LLM没有生成新报告，我们沿用上一轮的报告
                logger.warning("No <report> found. Report will not be updated for the next round.")

            # 5.2 检查是否终止
            if answer_content:
                prediction = answer_content
                termination = 'answer found'
                if terminate_flag:
                    termination = 'terminate with answer'
                await emit({
                    "type": "final",
                    "round": round_num,
                    "answer": prediction,
                    "report": research_round.current_report,
                    "termination": termination,
                })
                logger.debug(f"Round {round_num}: LLM provided <answer>. Terminating loop.")
                break
            if terminate_flag:
                if terminate_reason:
                    prediction = terminate_reason
                else:
                    prediction = research_round.current_report.strip()
                termination = 'terminated by llm'
                await emit({
                    "type": "final",
                    "round": round_num,
                    "answer": prediction,
                    "report": research_round.current_report,
                    "termination": termination,
                })
                logger.debug(f"Round {round_num}: LLM signaled <terminate>. Final response prepared.")
                break
            if 'is_last_call' in locals() and is_last_call:
                fallback_report = research_round.current_report.strip()
                fallback_source = fallback_report or terminate_reason
                prediction = fallback_source if fallback_source else research_round.last_observation
                termination = 'finalized without answer tag'
                await emit({
                    "type": "final",
                    "round": round_num,
                    "answer": prediction,
                    "report": research_round.current_report,
                    "termination": termination,
                })
                logger.warning("Last LLM call did not return <answer> or <terminate>. Promoting accumulated content as final answer.")
                break

            # 5.3 执行 Action (A_i)
            if action_content:
                try:
                    logger.debug(f"Round {round_num}: Executing tool...")
                    tool_response_str = await self.custom_call_tool(action_content)

                    # 将工具响应 O_i 存储，用于下一轮 s_{t+1}
                    research_round.last_observation = tool_response_str
                    await emit({
                        "type": "tool",
                        "round": round_num,
                        "tool_call": action_content,
                        "observation": tool_response_str,
                    })

                    # 记录工具响应到完整日志
                    tool_obs_msg = f"{OBS_START}\n{tool_response_str}\n{OBS_END}"
                    full_trajectory_log.append({"role": "user", "content": tool_obs_msg})

                    logger.debug(f"Round {round_num}: Tool execution completed.")

                except Exception as e:
                    logger.error(f"Error calling tool: {e}")
                    error_str = f"Error executing tool: {e}"
                    research_round.last_observation = error_str
                    await emit({
                        "type": "tool_error",
                        "round": round_num,
                        "tool_call": action_content,
                        "observation": error_str,
                    })
                    full_trajectory_log.append({"role": "user", "content": f"{OBS_START}\n{error_str}\n{OBS_END}"})
            else:
                # LLM 既没有回答也没有调用工具
                logger.warning("LLM did not produce <answer> or <tool_call>. Forcing answer generation...")

                # 强制生成答案
                force_answer_msgs = current_context + [
                    {"role": "user", "content": (
                        "You did not provide a valid response format. "
                        "Based on your current report and the information gathered so far, "
                        "please provide the final answer to the original question. "
                        "Use the three-part format: <plan>...</plan> <report>...</report> <answer>...</answer>"
                    )}
                ]

                try:
                    forced_content = await self.call_server(force_answer_msgs)
                    forced_parsed = self.parse_output(forced_content)

                    if forced_parsed["answer"]:
                        prediction = forced_parsed["answer"]
                        termination = "answer (forced)"
                        await emit({
                            "type": "final",
                            "round": round_num,
                            "answer": prediction,
                            "report": research_round.current_report,
                            "termination": termination,
                        })
                        full_trajectory_log.append({"role": "assistant", "content": forced_content})
                        break
                    else:
                        logger.error("Failed to force answer generation")
                        prediction = "No answer found (format error after retry)."
                        termination = "format error"
                        break
                except Exception as e:
                    logger.error(f"Error during forced answer generation: {e}")
                    prediction = "No answer found (format error)."
                    termination = "format error"
                    break

            # 5.4 Token 限制检查
            token_count = self.count_tokens(request_msgs if 'request_msgs' in locals() else current_context)
            logger.debug(f"Round {round_num} context token count: {token_count}")
            if token_count > self.max_input_tokens:
                logger.warning(f"Token quantity exceeds the limit: {token_count}")
                # 强制生成答案
                force_answer_msgs = current_context + [
                    {"role": "user", "content": "You have now reached the maximum context length. "
                                                "Stop making tool calls. Based on your research report, "
                                                "provide the final answer in the three-part format: "
                                                "<plan>...</plan> <report>...</report> <answer>...</answer>"}
                ]
                content = await self.call_server(force_answer_msgs)
                parsed = self.parse_output(content)
                prediction = parsed["answer"] if parsed["answer"] else "No answer found (token limit)."
                termination = 'token limit reached'
                await emit({
                    "type": "final",
                    "round": round_num,
                    "answer": prediction,
                    "report": research_round.current_report,
                    "termination": termination,
                })
                full_trajectory_log.append({"role": "assistant", "content": content})
                break

        # 循环结束后的收尾
        if not prediction:
            fallback_report = research_round.current_report.strip()
            if fallback_report:
                prediction = fallback_report
                if not termination:
                    termination = 'report fallback'
            elif num_llm_calls_available == 0:
                prediction = 'No answer found (exceeded available LLM calls).'
                termination = 'exceed available llm calls'
            else:
                prediction = 'No answer found.'
                termination = 'answer not found'

        # 返回最终结果
        result = {
            "question": question,
            "prediction": prediction,
            "report": research_round.current_report,
            "termination": termination,
            "trajectory": full_trajectory_log,
        }

        await emit({
            "type": "status",
            "status": termination or 'completed',
            "answer": prediction,
            "report": research_round.current_report,
        })

        return result


async def main():
    # 1. 定义你的 LLM 配置
    llm_config = {
        "model": LLM_MODEL_NAME,
        "generate_cfg": {
            'max_input_tokens': 32000,
            "temperature": 0.6,
            "top_p": 0.95,
            "presence_penalty": 1.1
        }
    }

    # 2. 实例化 agent
    agent = WebResearcherAgent(
        llm_config=llm_config,
        function_list=['search', 'python']
    )

    question = '刘翔破纪录时候是多少岁?'
    gt = "23岁差2天"

    # 4. 运行完整的并行研究与总结
    final_result = await agent.run(question)

    # 5. 打印结果
    print(f"final_result: {final_result}")


if __name__ == "__main__":
    asyncio.run(main())
