import xml.etree.ElementTree as ET
from datetime import datetime
from enum import StrEnum
from typing import Annotated

from pydantic import BaseModel, Field
from tongcell.core import ActorId

from judy_core.configs import JudyCompletionOption, LitellmCompletionOption
from judy_core.utils.common import truncate_text
from judy_core.utils.tool import ToolDefinition

TASK_PROMPT_TEMPLATE = """\
# 任务
<task>
{task_desc}
</task>
"""


class Task(BaseModel):
    task_desc: str
    env_notes: str | None = None
    verifier_mandatory_checklist: dict[str, str] | None = None

    def to_prompt(self) -> str:
        content = TASK_PROMPT_TEMPLATE.format(task_desc=self.task_desc)
        return content


class TodoStatus(StrEnum):
    PENDING = "PENDING"
    # IN_PROGRESS = "IN_PROGRESS"
    COMPLETED = "COMPLETED"


TeamMemberName = str


class TeamMemberId(BaseModel):
    kind: str
    name: TeamMemberName

    @classmethod
    def from_actor_id(cls, actor_id: ActorId):
        return cls(kind=actor_id.kind, name=actor_id.name)


class Todo(BaseModel):
    status: TodoStatus
    desc: str

    def to_xml_element(self) -> ET.Element:
        elem = ET.Element("todo")
        elem.text = self.desc
        return elem


class Plan(BaseModel):
    task_finished: Annotated[
        bool, Field(description="任务完成标记。在无需进一步规划时设置为为true。")
    ]
    todos: list[Todo]

    def to_prompt(self) -> str:
        plan_elem = ET.Element("plan")
        for i, todo in enumerate(self.todos):
            todo_element = todo.to_xml_element()
            todo_element.attrib["index"] = str(i + 1)
            todo_element.attrib["status"] = str(todo.status)
            plan_elem.append(todo_element)

        ET.indent(plan_elem, space=" ", level=0)
        return ET.tostring(plan_elem, encoding="unicode")

    def str_for_human(self) -> str:
        segs = [f"task_finished: {self.task_finished}"]
        for i, todo in enumerate(self.todos):
            todo_desc_lines = todo.desc.split("\n")
            todo_desc = "\n   ".join(todo_desc_lines)
            segs.append(f"{i + 1}. [{todo.status}] {todo_desc}")
        return "\n".join(segs)


class ExecutionEndReason(StrEnum):
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
    HAVE_SUGGESTIONS = "HAVE_SUGGESTIONS"
    STEPS_EXHAUSTED = "STEPS_EXHAUSTED"

    def str_for_model(self) -> str:
        if self == ExecutionEndReason.COMPLETED:
            return "成功"
        elif self == ExecutionEndReason.FAILED:
            return "失败"
        elif self == ExecutionEndReason.HAVE_SUGGESTIONS:
            return "提建议"
        elif self == ExecutionEndReason.STEPS_EXHAUSTED:
            return "步数耗尽"
        return "未知"


class ExecutorResult(BaseModel):
    member_id: TeamMemberId
    end_reason: ExecutionEndReason

    plan: Plan
    todo_index: int
    extra_finished_todo_indexes: list[int]
    report_content: str

    step_count: int
    tool_call_count: int


class VerificationRequest(BaseModel):
    task: Task
    task_execution_summary: str


class VerificationResultStatus(StrEnum):
    PASSED = "PASSED"
    FAILED = "FAILED"
    STEPS_EXHAUSTED = "STEPS_EXHAUSTED"


class VerificationResult(BaseModel):
    member_id: TeamMemberId
    status: VerificationResultStatus
    report_content: str


def executor_result_for_model(executor_result: ExecutorResult) -> ET.Element:
    executor_result_elem = ET.Element("executor_result")
    executor_result_elem.attrib["member_name"] = executor_result.member_id.name
    executor_result_elem.attrib["end_reason"] = executor_result.end_reason
    executor_result_elem.attrib["step_count"] = str(executor_result.step_count)
    executor_result_elem.attrib["tool_call_count"] = str(
        executor_result.tool_call_count
    )

    todo = executor_result.plan.todos[executor_result.todo_index]
    todo_element = todo.to_xml_element()
    executor_result_elem.append(todo_element)
    for index in executor_result.extra_finished_todo_indexes:
        extra_todo = executor_result.plan.todos[index]
        todo_element = extra_todo.to_xml_element()
        executor_result_elem.append(todo_element)

    REPORT_CONTENT_MAX_LEN = 2000
    report_truncation_result = truncate_text(
        executor_result.report_content, REPORT_CONTENT_MAX_LEN
    )
    report_elem = ET.Element("report_content")
    report_elem.attrib["length"] = str(len(executor_result.report_content))
    report_elem.text = report_truncation_result.text

    executor_result_elem.append(report_elem)
    return executor_result_elem


def verification_result_for_model(
    verification_result: VerificationResult,
) -> ET.Element:
    elem = ET.Element("verification_result")
    elem.attrib["member_name"] = verification_result.member_id.name
    elem.attrib["status"] = verification_result.status
    elem.text = verification_result.report_content
    return elem


type TeamMemberResult = ExecutorResult | VerificationResult


class TaskHistory(BaseModel):
    team_member_results: list[TeamMemberResult]

    def to_prompt(self):
        task_history_elem = ET.Element("task_history")
        for team_member_result in self.team_member_results:
            if isinstance(team_member_result, ExecutorResult):
                elem = executor_result_for_model(team_member_result)
            else:
                elem = verification_result_for_model(team_member_result)
            task_history_elem.append(elem)
        ET.indent(task_history_elem, space=" ", level=0)
        return ET.tostring(task_history_elem, encoding="unicode")


class Subtask(BaseModel):
    task: Task
    task_history: TaskHistory
    plan: Plan
    todo_index: int


BASIC_ENV_INFO_PROMPT_TEMPLATE = """\
# 一些可能有用的环境信息
这是系统检测到的一些环境信息，并不全面，仅供参考，更详细的需要自行探索。
<basic_env_info>
{basic_env_info}
</basic_env_info>
"""


class TaskConfig(BaseModel):
    space_dir: str
    task_dir: str | None = None
    basic_env_info: str | None = None
    disable_verifier: bool = False

    def to_prompt(self, agent_dir: str | None = None) -> str:
        segs = ["# 目录"]
        if self.task_dir:
            segs.append(
                "任务目录（默认工作目录）: `{task_dir}`".format(task_dir=self.task_dir)
            )
        segs.append("团队空间目录: `{space_dir}`".format(space_dir=self.space_dir))
        if agent_dir:
            segs.append(
                "你的个人目录: `{space_dir}/{agent_dir}`".format(
                    space_dir=self.space_dir, agent_dir=agent_dir
                )
            )

        segs.append("")

        if self.basic_env_info:
            segs.append(
                BASIC_ENV_INFO_PROMPT_TEMPLATE.format(
                    basic_env_info=self.basic_env_info
                )
            )
        content = "\n".join(segs)
        return content


class GetTaskConfig(BaseModel):
    pass


class GetModelConfig(BaseModel):
    pass


class ModelConfig(BaseModel):
    model: str
    tools_definition: list[ToolDefinition]
    judy_completion_option: JudyCompletionOption = Field(
        default_factory=JudyCompletionOption
    )
    litellm_completion_option: LitellmCompletionOption = Field(
        default_factory=LitellmCompletionOption
    )


class MemberQuestion(BaseModel):
    text: str


class MemberAnswer(BaseModel):
    text: str


class GetTaskState(BaseModel):
    pass


class TaskState(BaseModel):
    task_start_ts: float

    def __init__(self):
        super().__init__(task_start_ts=datetime.now().timestamp())

    def elapsed_time_sec(self) -> float:
        return datetime.now().timestamp() - self.task_start_ts


class ShellCommandCompletedNotification(BaseModel):
    text: str


class GetEnvironmentNotification(BaseModel):
    pass


type EnvironmentNotification = ShellCommandCompletedNotification


class EnvironmentNotificationList(BaseModel):
    notifications: Annotated[list[EnvironmentNotification], Field()] = []
