# -*- coding: utf-8 -*-
"""Verification manager for UCAgent stage execution."""

import copy
import os
import time
import traceback
import random
from collections import OrderedDict
from typing import Optional, Callable, Dict, Any

from langchain_core.callbacks import (
    CallbackManagerForToolRun,
)
from langchain_core.tools.base import ArgsSchema
from pydantic import BaseModel, Field

import ucagent.util.functions as fc
from ucagent.checkers import UnityChipCheckerTestFree
from ucagent.stage.vstage import get_root_stage
from ucagent.tools.uctool import UCTool, EmptyArgs
from ucagent.util.functions import make_llm_tool_ret
from ucagent.util.log import info, warning
from ucagent.stage.llm_suggestion.base_suggestion import get_llm_check_instance
from ucagent.tools.skill import _list_skills, list_skills_in_format


class ManagerTool(UCTool):
    # custom vars
    function: Callable = None
    args_schema: Optional[ArgsSchema] = EmptyArgs

    def _run(self, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        return self.function()

    def set_function(self, func):
        self.function = func
        return self


class ArgApproveStagePass(BaseModel):
    pass_or_not: bool = Field(
        default=True,
        description="Indicates whether to pass or not current stage task, True means pass and False means fail. Default is True."
    )


class ApproveStagePass(ManagerTool):
    """Approve the current stage as passed."""
    name: str = "ApproveStagePass"
    args_schema: Optional[ArgsSchema] = ArgApproveStagePass
    description: str = (
        "Approve the current stage as passed. \n"
        "This tool is used when you have verified that the current stage has been properly completed. \n"
    )
    def _run(self, pass_or_not: bool = True, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        return self.function(pass_or_not)


class ToolStatus(ManagerTool):
    """List current missoin status."""
    name: str = "Status"
    description: str = (
        "Returns the current status of your mission."
    )


class ToolCurrentTips(ManagerTool):
    """Get tips for the current task."""
    name: str = "CurrentTips"
    description: str = (
        "Returns the tips for the current task."
    )


class ToolDetail(ManagerTool):
    """Get current missoin detials."""
    name: str = "Detail"
    description: str = (
        "Returns the detail info of your mission, including all stages and their details. \n"
    )


class ToolKillCheck(ManagerTool):
    """Kill the current check process."""
    name: str = "KillCheck"
    description: str = (
        "Kill the current check process. \n"
        "This tool is only used when the tool 'Check' is long time running or get stuck. \n"
    )


class ToolStageJournal(ManagerTool):
    """Get the journal of the current stage."""
    name: str = "StageJournal"
    description: str = (
        "Get the journal of the current stage. "
    )


class ToolAllStageJournal(ManagerTool):
    """get the journal of the all stages."""
    name: str = "AllStageJournal"
    description: str = (
        "Get the journal of all stages. \n"
        "This tool is used to when continue a previous mission or the LLM context is compressed and other similar situations. "
    )


class ArgSetCurrentStageJournal(BaseModel):
    journal: str = Field(
        description="The journal content to set for the current stage. Cannot be empty."
    )


class ToolSetCurrentStageJournal(ManagerTool):
    """set the journal of the current stage."""
    name: str = "SetCurrentStageJournal"
    description: str = (
        "Set the journal of the current stage. \n"
        "This tool is used to record important information during the current stage. When completing the stage, the journal should be set. \n"
        "The journal content should be concise and clear and only the necessary information should be included.\n"
        "eg: - What you have done in this stage.\n"
        "    - What problems you have encountered and how you solved them.\n"
        "    - Experience or lessons learned during this stage.\n"
        "    - Files and its comments you have created or modified in this stage.\n"
        "    - Things to note when re-continuing this stage in the future.\n"
    )
    args_schema: Optional[ArgsSchema] = ArgSetCurrentStageJournal

    def _run(self, journal: str = "",
             run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        if not journal:
            return "Journal content cannot be empty."
        return self.function(journal)

class ArgSkillUsage(BaseModel):
    skill_usage: Dict[str, Any] = Field(
        description="The skill usage to set for the current stage. Cannot be empty."
    )

class ToolSetSkillUsage(ManagerTool):
    """Check and set the skill usage of the current stage."""
    name: str = "SetSkillUsage"
    description: str = (
        "Check the usage of the skills and set journal of the current stage. \n"
        "Analyze the conversation history and check usage of the skills specified in skill_list (if skills beyond the specified list were also used, analyze them as well).\n"
        "For each skill, analyze the following aspects:\n"
        "1. **list**: Whether the name and description of skill was listed in histoty context\n"
        "2. **read**: Whether the SKILL.md of skill was read by using tool `ReadTextFile`\n"
        "3. **use**: Whether completion of the current stage task followed the method steps in SKILL.md, or executed any specified code in that file\n"
        "**Returned dictionary format example**:\n"
        "{\n"
        "  'skill_name_1': {'list': True, 'read': True, 'use': False},\n"
        "  'skill_name_2': {'list': True, 'read': False, 'use': False}\n"
        "}\n"
    )
    args_schema: Optional[ArgsSchema] = ArgSkillUsage

    def _run(self, skill_usage: Dict[str, Any] = None, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        if not skill_usage:
            return "Skill usage content cannot be empty, use tool `ToolSetSkillUsage` to check the skill usage and set the skill usage content."
        return self.function(skill_usage)


class ArgStageDiff(BaseModel):
    target_file: str = Field(".", description="The target file or path to get diff, default is current workspace directory.")
    show_detail: bool = Field(False, description="Whether to show detailed diff output, default is False.")
    start_line: int = Field(1, description="The starting line number for diff output, default is 1")
    line_count: int = Field(-1, description="The number of lines to show in the diff output, default is -1 (show all lines)")


class ToolStageDiff(ManagerTool):
    """Retrieve the differences between the current file and the file at the last `StageCommit`."""
    name: str = "StageDiff"
    description: str = (
        "Get the differences between the current file and the file at the last `StageCommit`. \n"
        "This tool helps you to identify changes made since the last commit in the current stage. \n"
        "Use this tool to review modifications before proceeding with further actions. \n"
    )
    args_schema: Optional[ArgsSchema] = ArgStageDiff
    def _run(self, target_file: str = ".",
             show_detail: bool = False,
             start_line: int = 1,
             line_count: int = -1,
             run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        return self.function(target_file, show_detail, start_line, line_count)


class ArgStageCommit(BaseModel):
    commit_message: str = Field(..., description="The commit message for the changes in the current stage.")


class ToolStageCommit(ManagerTool):
    """Commit current stage changes."""
    name: str = "StageCommit"
    description: str = (
        "Commit current stage changes. \n"
        "This tool records your progress and changes made during the current stage. \n"
        "Use this tool to ensure that all modifications are saved before proceeding with further actions. \n"
        "When called this tool, all changes in the current stage will be committed with the provided commit message and "
        "you cannot undo this action. \n"
    )
    args_schema: Optional[ArgsSchema] = ArgStageCommit
    def _run(self, commit_message: str = "",
             run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        if not commit_message:
            return "Commit message cannot be empty."
        return self.function(commit_message)


class ArgStdCheck(BaseModel):
    lines: int = Field(
        default=-1,
        description="lines to read, -1 means read all"
    )


class ToolStdCheck(ManagerTool):
    """get the standard output of the current check process."""
    name: str = "StdCheck"
    description: str = (
        "Get the standard output of the current check process. \n"
        "This tool is only used to get the output of the runnig tool 'Check'. \n"
        "You can specify the number of lines to read, -1 means read all lines. \n"
    )
    args_schema: Optional[ArgsSchema] = ArgStdCheck

    def _run(self, lines: int = -1, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        return self.function(lines)


class ArgCheck(BaseModel):
    target: str = Field(
        default="",
        description=(
            "Target test cases to run, supports pytest-style arguments for precise test selection. "
            "Examples:\n"
            "• '' (empty): Run all test cases in the test directory\n"
            "• 'test_file.py': Run all tests in a specific file\n"
            "• 'test_file.py::test_function': Run a specific test function\n"
            "• 'test_file.py::TestClass::test_method': Run a specific test method in a class\n"
            "• '-k pattern': Run tests matching the given pattern\n"
            "• '-m marker': Run tests with specific markers\n"
        )
    )
    timeout: int = Field(
        default=0,
        description="Timeout for the test run in seconds. Zero means use default cfg.call_time_out."
    )
    return_line_coverage: bool = Field(
        default=False,
        description="Whether to return line coverage information in the test results."
    )


class ArgsDoCheck(BaseModel):
    timeout: int = Field(
        default=0,
        description="Timeout for Check/Complete tools. Zero means use default cfg.call_time_out."
    )
    ex_args: str = Field(
        default="",
        description="Args passed to the stage checkers. According to the stage help info to pass the corrent value."
    )


class ToolRunTestCases(ManagerTool):
    """Run test cases in current workspace."""
    name: str = "RunTestCases"
    description: str = (
        "This tool is used to execute the test cases in the workspace. "
        "Returns the result of the test execution. You should call this tool after you have implemented or modified the DUT or test cases. "
        "Current test directory is set to the '{TEST_DIR}',  the file path you passed should be relative to this directory."
    )
    args_schema: Optional[ArgsSchema] = ArgCheck

    def _run(self, target="", timeout=0, return_line_coverage=False,
             run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        try:
            if timeout <= 0:
                timeout = self.get_call_time_out()
            return self.function(target, timeout, return_line_coverage)
        except Exception as e:
            traceback.print_exc()
            error_msg = f"Test execution failed: {str(e)}"
            info(error_msg)
            return error_msg


class ToolDoCheck(ManagerTool):
    """Advanced validation tool for stage requirements and implementation quality."""
    name: str = "Check"
    description: str = (
        "Perform comprehensive validation of your current stage's implementation against requirements.\n"
        "The tool provides detailed feedback."
    )
    args_schema: Optional[ArgsSchema] = ArgsDoCheck

    def _run(self, timeout=0, ex_args="", run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        """
        Execute stage validation with enhanced error handling and reporting.
        
        Args:
            target: Test target specification (pytest format)
            ex_args: String args passed to checkers
            run_manager: Callback manager for tool execution
            
        Returns:
            str: Comprehensive validation report in JSON format
        """
        try:
            if timeout <= 0:
                timeout = self.get_call_time_out()
            return self.function(timeout, ex_args)
        except Exception as e:
            traceback.print_exc()
            error_msg = f"Validation failed: {str(e)}"
            info(error_msg)
            return make_llm_tool_ret({
                "check_pass": False,
                "check_info": error_msg
            })


class ToolDoComplete(ManagerTool):
    """Tell the manager that you have completed the current stage."""
    name: str = "Complete"
    description: str = (
        "Perform comprehensive validation of your current stage's implementation against requirements and mark the stage as complete if all checks pass.\n"
        "The tool provides detailed feedback (Different from tool 'Check': if all checks pass, the stage is marked as complete and the manager advances to the next stage).\n\n"
    )
    args_schema: Optional[ArgsSchema] = ArgsDoCheck

    def _run(self, timeout=0, ex_args="", run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        try:
            if timeout <= 0:
                timeout = self.get_call_time_out()
            return self.function(timeout, ex_args)
        except Exception as e:
            traceback.print_exc()
            error_msg = f"Completion failed: {str(e)}"
            info(error_msg)
            return error_msg


class ArgToolGoToStage(BaseModel):
    index: int = Field(
        default=-1,
        description="Stage index to go to. "
    )


class ToolGoToStage(ManagerTool):
    """Go to a specific stage by index."""
    name: str = "GoToStage"
    description: str = (
        "Go to a specific stage by index. Only those stages that have been reached can be selected. \n"
        "Stage is reached means that all checks in the stage have been passed. \n"
        "This tool is used when you want refine your previous work, or want to go back to a previous stage. \n"
        "Returns the result of the operation."
    )
    args_schema: Optional[ArgsSchema] = ArgToolGoToStage

    def _run(self, index: int = -1, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        return self.function(index)


class ToolDoExit(ManagerTool):
    """Exit the agent and end the mission after all stages are completed."""
    name: str = "Exit"
    description: str = (
        "Exit the agent and end the mission after all stages are completed. \n"
        "This tool is used when you have completed all stages and want to exit the agent. \n"
        "Returns a message indicating the exit status."
    )


class StageManager(object):
    def __init__(
            self, workspace, cfg, agent, tool_read_text, ucagent_info: dict,
            force_stage_index=0,
            force_todo=False,
            todo_panel=None,
            stage_skip_list=None,
            stage_unskip_list=None,
            tool_inspect_file=None,
            reference_files=None,
    ):
        """
        Initialize the StageManager with an empty list of stages.
        """
        self.cfg = cfg
        self.data = {}
        self.workspace = workspace
        self.force_todo = force_todo
        self.todo_panel = todo_panel
        self.free_pytest_run = UnityChipCheckerTestFree("", cfg.tools.RunTestCases.test_dir, "").set_workspace(workspace)
        self.agent = agent
        self.tool_read_text = tool_read_text
        self.ucagent_info = ucagent_info
        self.force_stage_index = force_stage_index
        self.stage_skip_list = stage_skip_list
        self.stage_unskip_list = stage_unskip_list
        self.tool_inspect_file = tool_inspect_file
        self.reference_files = reference_files

    def init_stage(self):
        from ucagent.stage import VerifyStage
        self.root_stage = get_root_stage(self.cfg, self.workspace, self.tool_read_text)
        self.stages = self.root_stage.get_substages()
        if self.reference_files:
            for si, flist in self.reference_files.items():
                if 0 <= si < len(self.stages):
                    info(f"Stage {si} try add reference files: {flist}")
                    self.stages[si].add_reference_files(flist)
                elif si == -1:
                    info(f"All stages try add reference files: {flist}")
                    for s in self.stages:
                        s.add_reference_files(flist)
                else:
                    warning(f"Invalid stage index {si} in reference_files, ignored.")
        self.mission = self.cfg.mission
        info(f"Initialized StageManager with {len(self.stages)} stages.")
        info("Stages:\n" + "\n".join([f"{i:2d}:   {stage.title()}{' (skipped)' if stage.is_skipped() else ''}" for i, stage in enumerate(self.stages)]))
        self.stage_index = min(max(0, self.force_stage_index), len(self.stages))
        for i in range(min(self.stage_index + 1, len(self.stages))):
            self.stages[i].set_reached(True)
        stages_info = self.ucagent_info.get("stages_info", {})

        for stage_idx_str, stage_info in stages_info.items():
            idx = int(stage_idx_str)
            if idx >= len(self.stages):
                continue
            stage: VerifyStage = self.stages[idx]
            stage.set_fail_count(stage_info.get("fail_count", 0))
            stage.set_time_prev_cost(stage_info.get("time_cost", 0.0))
            stage.set_reached(stage_info.get("reached", stage.is_reached()))
            stage.set_skip(stage_info.get("is_skipped", stage.is_skipped()))
            stage.is_complete = stage_info.get("is_completed", stage.is_completed())
            stage.set_reference_file_status(stage_info.get("task", {}).get("reference_files", {}))
            if "meta_data" in stage_info:
                stage.meta_data = copy.deepcopy(stage_info["meta_data"])
        self._go_skip_stage()
        for s in self.stages:
            s.set_stage_manager(self)
        if self.stage_index < len(self.stages):
            self.stages[self.stage_index].on_init()
        self.last_check_info = {}
        self.all_completed = bool(self.ucagent_info.get("all_completed", False))
        if self.stage_skip_list:
            for si in self.stage_skip_list:
                self.skip_stage(si)
                info(f"Stage {si} is set to be skipped.")
        if self.stage_unskip_list:
            for sui in self.stage_unskip_list:
                self.unskip_stage(sui)
                info(f"Stage {sui} is set to be unskipped.")
        self._refresh_all_completed()
        info("Current stage index is " + str(self.stage_index) + ".")
        self.time_begin = self.ucagent_info.get("time_begin", time.time())
        self.time_end = self.ucagent_info.get("time_end", None)
        self.llm_fail_suggestion = get_llm_check_instance(
            self.cfg.vmanager.llm_suggestion.check_fail_refinement,
            self,
            self.tool_inspect_file
        )
        self.llm_pass_suggestion = get_llm_check_instance(
            self.cfg.vmanager.llm_suggestion.check_pass_refinement,
            self,
            self.tool_inspect_file + [ApproveStagePass().set_function(self.tool_stage_approve),
                                      ToolStageDiff().set_function(self.tool_stage_diff),
                                      ToolStageCommit().set_function(self.tool_stage_commit)]
        )
        if self.stage_index < len(self.stages):
            self.stages[self.stage_index].hist_init()

    def is_break(self):
        return self.agent.is_break()

    def tool_stage_approve(self, pass_or_not: bool = True) -> str:
        vstage = self.get_current_stage()
        if vstage is None:
            return "No current stage available."
        vstage.set_approved(pass_or_not)
        return f"Stage '{vstage.name}' approved status set to {pass_or_not}."

    def tool_stage_diff(self, target_file, show_detail, start_line, line_count) -> str:
        vstage = self.get_current_stage()
        if vstage is None:
            return "No current stage available."
        return vstage.hist_diff(target_file, show_detail, start_line, line_count)

    def tool_stage_commit(self, commit_message) -> str:
        vstage = self.get_current_stage()
        if vstage is None:
            return "No current stage available."
        vstage.hist_commit(commit_message)
        return f"Stage '{vstage.name}' changes committed."

    def gen_fail_suggestion(self, error_msg) -> str:
        stage = self.get_current_stage()
        if stage is None:
            return error_msg
        try:
            fail_suggestion = self.gen_llm_suggestion(
                error_msg,
                stage,
                self.llm_fail_suggestion,
                self.stage_need_llm_fail_suggestion,
            )
            stage.meta_set_llm_fail_suggestion(fail_suggestion)
            return fail_suggestion
        except Exception as e:
            traceback.print_exc()
            warning(f"Generate fail suggestion failed: {str(e)}")
            return error_msg

    def is_llm_suggestion_needed(self, stage_name, suggestion_ins, need_llm_suggestion) -> bool:
        if not suggestion_ins:
            return False
        if need_llm_suggestion is not None:
            return need_llm_suggestion
        cfg = suggestion_ins.get_cfg()
        bypass_stages = cfg.get("bypass_stages", [])
        if bypass_stages:
            if stage_name in bypass_stages:
                return False
        target_stages = cfg.get("target_stages", [])
        if target_stages:
            if stage_name not in target_stages:
                return False
        return cfg.get("default_apply_all_stages", True)

    def stage_need_llm_pass_suggestion(self, stage) -> bool:
        if stage is None:
            return False
        return self.is_llm_suggestion_needed(
            stage.name,
            self.llm_pass_suggestion,
            stage.need_pass_llm_suggestion,
        )

    def stage_need_llm_fail_suggestion(self, stage) -> bool:
        if stage is None:
            return False
        return self.is_llm_suggestion_needed(
            stage.name,
            self.llm_fail_suggestion,
            stage.need_fail_llm_suggestion,
        )

    def gen_pass_suggestion(self, raw_msg) -> str:
        stage = self.get_current_stage()
        if stage is None:
            return raw_msg
        try:
            pass_suggestion = self.gen_llm_suggestion(
                raw_msg,
                stage,
                self.llm_pass_suggestion,
                self.stage_need_llm_pass_suggestion,
                pre_llm_cb = lambda: stage.set_approved(False),
            )
            stage.meta_set_llm_pass_suggestion(pass_suggestion)
            return pass_suggestion
        except Exception as e:
            traceback.print_exc()
            warning(f"Generate pass suggestion failed: {str(e)}")
            warning("Set stage as approved due to exception in generating pass suggestion.")
            stage.set_approved(True)
            return raw_msg

    def gen_llm_suggestion(self, raw_msg, stage,
                           suggestion_instance,
                           stage_need_llm_suggestion_fc,
                           pre_llm_cb=None) -> str:
        assert stage is not None, "stage can not be None"
        if stage_need_llm_suggestion_fc(stage) is False:
            return raw_msg
        if callable(pre_llm_cb):
            pre_llm_cb()
        stage.set_force_unactive(True)
        try:
            return suggestion_instance.suggest([
                    stage.task_info(),
                    raw_msg],
                    stage)
        finally:
            stage.set_force_unactive(False)

    def get_time_cost(self):
        if self.time_end is None:
            return time.time() - self.time_begin
        return self.time_end - self.time_begin

    def _compute_all_completed(self):
        if not self.stages:
            return True
        if self.stage_index >= len(self.stages):
            return True
        return all(stage.is_skipped() or stage.is_completed() for stage in self.stages)

    def _refresh_all_completed(self):
        self.all_completed = self._compute_all_completed()
        return self.all_completed

    def attach_todo_summary(self, data):
        assert isinstance(data, str), "the target data type of attach_todo_summary must be str"
        if not self.force_todo:
            return data
        if not self.todo_panel:
            return data
        return data + self.todo_panel._summary()

    def set_data(self, key, value):
        self.data[key] = value

    def get_data(self, key, default=None):
        return self.data.get(key, default)

    def new_tools(self):
        """
        Create and return a list of tools for the current stage.
        """
        tools = [
            ToolCurrentTips().set_function(self.tool_current_tips),
            ToolDetail().set_function(self.tool_detail),
            ToolStatus().set_function(self.tool_status),
            ToolRunTestCases().set_function(self.tool_run_test_cases).render_desc({"TEST_DIR": self.free_pytest_run.test_dir}),
            ToolDoCheck().set_function(self.tool_check),
            ToolKillCheck().set_function(self.tool_kill_check),
            ToolStdCheck().set_function(self.tool_std_check),
            ToolDoComplete().set_function(self.tool_complete),
            ToolGoToStage().set_function(self.tool_go_to_stage),
            ToolDoExit().set_function(self.tool_exit),
            ToolStageJournal().set_function(self.tool_get_current_journal),
            ToolAllStageJournal().set_function(self.tool_get_all_journal),
            ToolSetCurrentStageJournal().set_function(self.tool_set_journal),
        ]
        if self.agent.cfg.skill.use_skill:
            tools.append(ToolSetSkillUsage().set_function(self.tool_set_skill_usage))
        return tools

    def get_current_tips(self):
        if self.stage_index >= len(self.stages):
            return "Your mission is completed. No more stages available. You can use `Exit` tool to exit the mission or `GoToStage` tool to go to a specific stage to review."
        cstage = self.stages[self.stage_index]
        tips = OrderedDict()
        tips["mission"] = self.mission.name
        tips["current_stage"] = OrderedDict({
            "index": self.stage_index,
            **cstage.detail(),
        })
        ref_files = []
        for k, v in cstage.reference_files.items():
            if v:
                continue
            ref_files.append(k)
        if ref_files:
            tips["notes"] = f"You need use tool: {self.tool_read_text.name} to read the reference files.\n"
        
        # list the skills needed to use in current stage
        skills_to_use = [skill_name for skill_name in cstage.skill_list]
        if skills_to_use:
            formatted_skill_list = list_skills_in_format(_list_skills(self.workspace), self.workspace, skills_to_use)
            tips["notes"] = tips.get("notes", "") + f"Firstly you must read the SKILL.md of the following skills to know how to complete current stage:\n{formatted_skill_list}\n"

        tips["process"] = f"{self.stage_index}/{len(self.stages)}"
        mession_tips = self.mission.get_value("prompt.tips")
        if mession_tips is not None:
            mession_tips = mession_tips.as_dict()
            current_tips = mession_tips.get("allways", [])
            random_tips = mession_tips.get("random", [])
            if random_tips:
                rindex = random.randint(0, len(random_tips) - 1)
                current_tips.append(random_tips[rindex])
            if current_tips:
                random.shuffle(current_tips)
                tips["tips"] = current_tips
        tips = make_llm_tool_ret(tips)
        return self.attach_todo_summary(tips)

    def detail(self):
        """
        Get the details of the current mission, including all stages and their details.
        """
        ret = OrderedDict()
        ret["mission"] = self.mission.name
        ret["stage_list"] = []
        for i, stage in enumerate(self.stages):
            ret["stage_list"].append(stage.detail())
            ret["stage_list"][-1]["index"] = i
        ret["current_stage_index"] = self.stage_index
        ret["current_stage_name"] = self.stages[self.stage_index].name if self.stage_index < len(self.stages) else None
        return ret

    def status(self):
        ret = OrderedDict()
        ret["mission"] = self.mission.name
        ret["all_completed"] = self._compute_all_completed()
        ret["stage_list"] = []
        for i, stage in enumerate(self.stages):
            ret["stage_list"].append({
                "index": i,
                "title": stage.title(),
                "reached": stage.is_reached(),
                "fail_count": stage.fail_count,
                "skill_list": list(stage.skill_list.keys()) if self.cfg.skill.use_skill else [],
                "is_skipped": stage.is_skipped(),
                "time_start": stage.get_time_start_str(),
                "time_end": stage.get_time_end_str(),
                "time_cost": stage.get_time_cost_str(),
                "is_completed": stage.is_completed(),
                "needs_human_check": stage.is_hmcheck_needed(),
                "need_fail_llm_suggestion": self.stage_need_llm_fail_suggestion(stage),
                "need_pass_llm_suggestion": self.stage_need_llm_pass_suggestion(stage),
            })
        ret["process"] = f"{self.stage_index}/{len(self.stages)}"
        cstage = self.stages[self.stage_index] if self.stage_index < len(self.stages) else None
        ret["current_task"] = "No stages available (Maybe mission is completed, you can use the `GoToStage` tool to go back to a previous stage if needed)"
        if cstage:
            ret["current_stage_index"] = self.stage_index
            ret["current_stage_name"] = cstage.name
            ret["current_task"] = cstage.task_info()
        ret["last_check_result"] = self.last_check_info
        return ret

    def get_current_stage(self):
        return self.get_stage(self.stage_index)

    def set_current_stage_journal(self, journal):
        stage = self.get_current_stage()
        if stage:
            stage.meta_set_journal(journal)
            return "Set journal success."
        return "No current stage available."

    def get_current_stage_journal(self):
        stage = self.get_current_stage()
        if stage:
            return stage.meta_get_journal()
        return "No current stage available."

    def get_all_stage_journal(self):
        journals = OrderedDict()
        for stage in self.stages:
            journals[stage.title()] = stage.meta_get_journal()
        return journals

    def set_current_stage_skill_usage(self, skill_usage: Dict[str, Any]):
        """set the skill usage of curretn stage or return feedback based on skill_usage"""
        current_stage = self.get_current_stage()
        if current_stage.skill_list:
            for skill_name in current_stage.skill_list:
                skill_root = fc.get_workspace_skill_root(self.workspace)
                skill_dir = fc.find_skill_dir_by_name(skill_root, skill_name)
                if not skill_dir:
                    raise ValueError(f"Skill '{skill_name}' is not found in workspace. ")
                if skill_name not in skill_usage:
                    return f"You must use skill '{skill_name}' in current stage, using tool `ListSkill` to list and use it."
                else:
                    skill_info = skill_usage[skill_name]
                    current_stage.set_usage_skill_list(skill_name, listed=skill_info.get("list", False), read=skill_info.get("read", False), used=skill_info.get("use", False))
                    [u,v,w] = current_stage.skill_list[skill_name]
                    if u and v and w:
                        continue
                    if not u:
                        return f"You must re-complete the stage by using tool `ListSkill` to list and use the skill {skill_name}."
                    if not v:
                        return f"You must re-complete the stage by using tool `ReadTextFile` to read the SKILL.md of skill {skill_name} and use it"
                    if not w:
                        return f"You must re-complete the stage by using the skill {skill_name} according to the method steps mentioned in its SKILL.md."
            current_stage.meta_set_skill_usage(skill_usage)
            return "All skills in skill_list have been used."     
        return "No skill need be used in current stage."

    def get_stage(self, index):
        if 0 <= index < len(self.stages):
            return self.stages[index]
        return None

    def go_to_stage(self, index):
        """
        Go to a specific stage by index.
        """
        success = False
        if 0 <= index < len(self.stages):
            if index == self.stage_index:
                msg = f"Already at stage {index}: {self.stages[index].name}."
                success = True
            elif self.stages[index].is_skipped():
                msg = f"Can not goto the skipped stage"
            elif self.stages[index].is_reached():
                self.stage_index = index
                msg = f"Changed to stage {index}: {self.stages[index].name} success."
                success = True
            else:
                msg = f"Stage {index} is not reached yet. Can only go to stages that have been reached. You can use tool `ToolStaus` to find all reached stages."
        else:
            msg = f"Invalid stage index: {index}. No change made."
        return {"message": msg, "success": success}

    def force_go_to_stage(self, index):
        """
        Force go to a specific stage by index, ignoring whether it is reached or not.
        This is used when initializing the StageManager with a specific stage index.
        """
        if 0 <= index < len(self.stages):
            self.stage_index = index
            return True
        return False

    def check(self, timeout, ex_args=""):
        if not self.stage_index < len(self.stages):
            return OrderedDict({
                "check_pass": False,
                "check_info": f"Stage index{self.stage_index} out of range. (Mission maybe completed, you can use the `GoToStage` tool to go back to a previous stage if needed)",
            })
        ck_pass, ck_info = self.stages[self.stage_index].do_check(**{"timeout": timeout, "ex_args": ex_args})
        ret_data = OrderedDict({
            "check_info": ck_info,
            "check_pass": ck_pass,
        })
        if not ck_pass:
            ret_data["action"] = "Please fix the issues reported in 'check_info.last_msg.error' according to the suggestions, and then use the `Check` tool again to re-validate your work."
        self.last_check_info = copy.deepcopy(ret_data)
        if ck_pass:
            ret_data["message"] = f"Congratulations! Stage {self.stage_index} checks passed successfully, you can use tool 'Complete' to finish this stage."
        else:
            return self.gen_fail_suggestion(ret_data)
        return ret_data

    def save_stage_info(self):
        all_completed = self._refresh_all_completed()
        info = self.agent.get_stat_info()
        info.update({
            "stage_index": self.stage_index,
            "all_completed": all_completed,
            "time_begin": self.time_begin,
            "time_end": self.time_end,
            "is_agent_exit": self.agent.is_exit(),
        })
        info["stages_info"] = {}
        for idx, stage in enumerate(self.stages):
            stage_info = stage.detail()
            stage_info["time_cost"] = stage.get_time_cost()
            stage_info["meta_data"] = stage.meta_data
            info["stages_info"][idx] = stage_info
        stage = self.get_current_stage()
        is_wait_human_check = False
        if stage:
            is_wait_human_check = stage.is_wait_human_check()
        info["is_wait_human_check"] = is_wait_human_check
        fc.save_ucagent_info(self.workspace, info)

    def next_stage(self):
        self.stage_index += 1
        self._go_skip_stage()
        self._refresh_all_completed()
        self.save_stage_info()
        return self.get_current_stage()

    def _go_skip_stage(self):
        if self.stage_index >= len(self.stages):
            return
        sk = 0
        while self.stages[self.stage_index].is_skipped():
            self.stage_index += 1
            sk += 1
            if self.stage_index >= len(self.stages):
                break
        if sk > 0:
            info(f"skipped {sk} stages, current stage index is now {self.stage_index}.")

    def skip_stage(self, index):
        if 0 <= index < len(self.stages):
            self.stages[index].set_skip(True)
            info(f"Stage '{self.stages[index].name}' is set to be skipped.")
            if index == self.stage_index:
                self.next_stage()
        else:
            warning(f"Invalid stage index: {index}, can not set skip.")

    def unskip_stage(self, index):
        if 0 <= index < len(self.stages):
            self.stages[index].set_skip(False)
            info(f"Stage '{self.stages[index].name}' is set to be unskipped.")
        else:
            warning(f"Invalid stage index: {index}, can not set unskip.")

    def _stage_complete(self, stage):
        stage.on_complete()
        if self.llm_fail_suggestion:
            self.llm_fail_suggestion.on_stage_complete(stage)
        if self.llm_pass_suggestion:
            self.llm_pass_suggestion.on_stage_complete(stage)

    def complete(self, timeout, ex_args=""):
        if self.stage_index >= len(self.stages):
            return {
                "complete": False,
                "message": ("No more stages to complete. You can review your work and use the `GoToStage` tool to go back to a previous stage if needed. "
                            "Or you can use the `Exit` tool to exit the mission."),
                "last_check_result": self.last_check_info,
            }
        ck_pass, ck_info = self.stages[self.stage_index].do_check(**{"timeout": timeout, "is_complete": True, "ex_args": ex_args})
        stage = self.stages[self.stage_index]
        if ck_pass:
            if stage.meta_get_journal() is None:
                return {"complete": False,
                        "error": "Please use tool 'SetCurrentStageJournal' to set the journal of this stage before completing it."}
        if ck_pass:
            llm_msg = self.gen_pass_suggestion(ck_info)
            ck_pass = stage.get_approved()
            if not ck_pass:
                if isinstance(llm_msg, str):
                    return {"complete": False, "error": "Stage Complete Fail:\n\n" + llm_msg}
                return {"complete": False, "error": "Stage Complete Fail:\n\n" + llm_msg}
        if ck_pass and stage.is_hmcheck_needed():
            hm_passed, ck_msg = stage.get_hmcheck_state()
            if hm_passed is None:
                self.agent._need_human = True
                return {"error": ("Now you have passed the self check of this stage, but human check is needed before completing this stage. "
                                  "Please give a bref introduction of your work to help the human reviewer understand your implementation. "
                                  "Then wait for human review and approval."),
                        "complete": False}
            elif hm_passed is False:
                self.agent._need_human = True
                return {"error": ("Human check did not approve your work for this stage. "
                                  "Please address the issues raised by the human reviewer and then use the `Complete` tool again to complete this stage."),
                         "human_review_msg": ck_msg,
                         "complete": False
                        }
            else:
                assert hm_passed is True, "hm_passed should be True here"
                info("Human check approved for stage " + stage.name)

        self.last_check_info = OrderedDict({
            "check_info": ck_info,
            "check_pass": ck_pass,
        })
        if ck_pass:
            message = f"Stage {self.stage_index} completed successfully. "
            self._stage_complete(self.stages[self.stage_index])
            self.next_stage()
            if self.all_completed:
                message = ("All stages completed successfully. "
                           "Now you should review your work to check if everything is correct and all the users needs are matched. "
                           "When you are confident that everything is fine, you can use the `Exit` tool to exit the mission. "
                           )
            else:
                message += f"Current stage index is now {self.stage_index}. Use `CurrentTips` tool to get your new task. "
                self.stages[self.stage_index].set_reached(True)
                self.stages[self.stage_index].on_init()
        else:
            message = f"Stage {self.stage_index} not completed. Please check the task requirements."
        ret = OrderedDict({
            "complete": ck_pass,
            "message": message,
            "last_check_result": self.last_check_info,
        })
        if not ck_pass:
            ret["action"] = "Please fix the issues reported in 'last_check_result.check_info.last_msg.error' according to the suggestions, and then use the `Complete` tool again to complete this stage."
            return self.gen_fail_suggestion(self.last_check_info)
        return ret

    def exit(self):
        """
        Exit the agent and end the mission after all stages are completed.
        """
        if self._refresh_all_completed():
            self.time_end = time.time()
            self.agent.exit()  # Exit the agent if all stages are completed
            self.save_stage_info()
            self.agent.try_exit_on_completion()
            ex_msg = ""
            if self.agent._exit_on_completion:
                ex_msg = " UCAgent has quit. The MCP server is shutting down — all MCP tools will become unavailable. You need to stop Now!"
            return {
                "exit": True,
                "message": "All stages completed. Exiting the mission." + ex_msg
            }
        return {
            "exit": False,
            "message": "Not all stages are completed yet. Please complete all stages before exiting."
        }

    def tool_set_journal(self, journal):
        """
        Set the journal of the current stage.
        This is used to when current stage is completed or the LLM context is compressed and other similar situations.
        The journal content should be concise and clear and only the necessary information should be included.
        """
        ret = make_llm_tool_ret(self.set_current_stage_journal(journal))
        info("ToolSetCurrentStageJournal:\n" + ret)
        return self.attach_todo_summary(ret)

    def tool_get_all_journal(self):
        ret = make_llm_tool_ret(self.get_all_stage_journal())
        info("ToolGetAllStageJournal:\n" + ret)
        return self.attach_todo_summary(ret)

    def tool_get_current_journal(self):
        ret = make_llm_tool_ret(self.get_current_stage_journal())
        info("ToolGetCurrentStageJournal:\n" + ret)
        return self.attach_todo_summary(ret)
    
    def tool_set_skill_usage(self, skill_usage: Dict[str, Any]):
        ret = make_llm_tool_ret(self.set_current_stage_skill_usage(skill_usage))
        info("ToolSetSkillUsage:\n" + ret)
        return self.attach_todo_summary(ret)

    def tool_detail(self):
        """
        Get the details of the current mission, including all stages and their details.
        """
        detail = make_llm_tool_ret(self.detail())
        info("ToolDetail:\n" + detail)
        return self.attach_todo_summary(detail)

    def tool_status(self):
        stat = make_llm_tool_ret(self.status())
        info("ToolStatus:\n" + stat)
        return self.attach_todo_summary(stat)

    def tool_go_to_stage(self, index):
        ret = make_llm_tool_ret(self.go_to_stage(index))
        info("ToolGoToStage:\n" + ret)
        return self.attach_todo_summary(ret)

    def tool_check(self, timeout, ex_args=""):
        ret = make_llm_tool_ret(self.check(timeout, ex_args))
        info("ToolCheck:\n" + ret)
        return self.attach_todo_summary(ret)

    def tool_exit(self):
        ret = make_llm_tool_ret(self.exit())
        info("ToolExit:\n" + ret)
        return ret

    def tool_complete(self, timeout, ex_args=""):
        ret = make_llm_tool_ret(self.complete(timeout, ex_args))
        info("ToolComplete:\n" + ret)
        return self.attach_todo_summary(ret)

    def tool_kill_check(self):
        """
        Kill the current check process.
        This is used when the tool 'Check' is long time running or get stuck.
        """
        if not self.stage_index < len(self.stages):
            return f"Stage index({self.stage_index}) out of range. (Maybe mission is completed, you can use the `GoToStage` tool to go back to a previous stage if needed)"
        stage = self.stages[self.stage_index]
        ret = stage.do_kill()
        info("KillCheck:\n" + ret)
        return ret

    def tool_std_check(self, lines=-1):
        """
        Get the standard output of the current check process.
        This tool is only used to get the output of the running tool 'Check'.
        You can specify the number of lines to read, -1 means read all lines.
        """
        if not self.stage_index < len(self.stages):
            return f"Stage index({self.stage_index}) out of range. (Maybe mission is completed, you can use the `GoToStage` tool to go back to a previous stage if needed)"
        stage = self.stages[self.stage_index]
        ret = stage.do_std(lines)
        info("StdCheck:\n" + ret)
        return ret

    def tool_current_tips(self):
        """
        Get the tips for the current task.
        This is used to provide guidance to the user on what to do next.
        """
        tips = self.get_current_tips()
        info("Tips:\n" + tips)
        return tips

    def tool_run_test_cases(self, pytest_args="", timeout=0, return_line_coverage=False, raw_return=False, detail=False):
        """
        Run test cases.
        This tool is used to execute the test cases in the workspace.
        """
        ret = self.free_pytest_run.do_check(pytest_args, timeout=timeout, return_line_coverage=return_line_coverage, detail=detail)
        if raw_return:
            return ret
        ret = make_llm_tool_ret(ret[1])
        info("RunTestCases:\n" + ret)
        return self.attach_todo_summary(ret)
