import asyncio
import json
import xml.etree.ElementTree as ET
from typing import TYPE_CHECKING, Annotated, Literal, override
from anyxml import AnyXml
from dais_sdk.types import UserMessage
from pydantic import BaseModel
from src.db import db_context
from src.services.tasks import SubtaskService
from src.db.models import tasks as tasks_models
from src.schemas.tasks import subtask as subtask_schemas
from src.schemas.tasks import runtime as task_runtime_schemas
from ..toolset_wrapper import builtin_tool, BuiltinToolset
from ...types import TaskError, TaskInterrupted, TaskWaitingAction, TaskFinished, is_agent_tool_metadata

if TYPE_CHECKING:
    from ...task import AgentTask


class SubtaskToolRespond(BaseModel):
    answer: Annotated[str,
                      "The answer to provide for a response-pending tool call."]
    call_id: Annotated[str,
                       "The call_id of the pending tool call, as returned in the subtask result."]

class SubtaskToolApprove(BaseModel):
    decision: Annotated[Literal["approved", "denied"],
                        "Whether to approve or deny the pending tool call."]
    call_id: Annotated[str,
                       "The call_id of the pending tool call, as returned in the subtask result."]

class NewSubtask(BaseModel):
    instruction: Annotated[str, "The initial instruction used to create a new subtask."]
    agent_id: int

class ContinueSubtask(BaseModel):
    subtask_id: Annotated[int,
                          "The ID of the subtask to continue, as returned in the subtask result."]
    agent_id: Annotated[int | None,
                        """
                        The id of target agent to execute the subtask.
                        Sometimes the selected agent may be deleted, can use this parameter to pass a new agent_id.
                        """] = None
    message: Annotated[str | list[SubtaskToolRespond | SubtaskToolApprove],
                       """
                       Either:
                       - A follow-up instruction when the subtask status is 'finished'
                       - Responses to pending tool calls when status is 'waiting_action'.
                         For each <tool_call> in the subtask result, pick the response type by pending_action:
                         - pending_action="respond" → SubtaskToolRespond (provide a text answer)
                         - pending_action="approve" → SubtaskToolApprove (approve or deny execution)
                       Match each response to its tool call via call_id.
                       """]

async def create_agent_task_from_subtask(subtask: tasks_models.Subtask) -> AgentTask:
    from ...task import AgentTask
    from ...context import AgentContext
    task_runtime = task_runtime_schemas.TaskRuntimeContext.from_subtask(subtask)
    ctx = await AgentContext.create(task_runtime)
    return AgentTask(ctx)

class OrchestrationToolset(BuiltinToolset):
    @property
    @override
    def name(self) -> str: return "Orchestration"

    @builtin_tool(validate=True)
    async def subtask(self, action: NewSubtask | ContinueSubtask) -> str:
        """
        Create a new subtask or continue an existing subtask.

        When to use:
            - To delegate a self-contained unit of work to a subtask
            - To isolate investigative or exploratory work whose intermediate steps (raw file contents, search output, dead-ends) would otherwise accumulate in the main context without contributing to it
            - To follow up on a completed subtask (e.g. ask for clarification or additional details)
            - To respond to a pending tool call from a running subtask (e.g. answering ask_user, approving or denying tool calls)

        When used in parallel:
            - Use 1 subtask when the goal is narrow and the context is clear. Use multiple subtasks concurrently when the task is broad, requires exploring multiple distinct aspects, or can be broken down into independent subtasks.
            - Quality over quantity, 5 subtasks maximum, but you should try to use the minimum number of subtasks necessary (usually just 1)
            - IMPORTANT: Never continue the same subtask in parallel. Multiple continue actions targeting the same subtask_id must be issued sequentially, not concurrently.

        On error:
            - If a subtask returns an error, retry by continuing with the *same* subtask_id rather than creating a new one.
              The subtask retains its accumulated context, so a continue is cheaper and preserves progress.
            - If the same subtask_id fails repeatedly (3+ consecutive errors), abandon it and open a new subtask.

        Preferred scenario:
            - To delegate directly to a specialized agent whose description matches the task at hand

        Consuming results:
            Once a subtask has investigated something, treat its findings as authoritative and do not re-investigate the same question inline.

        Returns:
            An XML string describing the current subtask result.
            The root element always has a `subtask_id` and `status` attribute.
            The `status` can be "finished", "waiting_action", "error", or "interrupted".

        Examples:
            When subtask finished with task progress:
                <subtask_result subtask_id="1" status="finished">
                    <progress>
                        <step status="completed">Research existing codebase</step>
                        <step status="completed">Implement feature</step>
                    </progress>
                    The task has been finished...
                </subtask_result>

            When subtask is waiting for a response or approval:
                <subtask_result subtask_id="1" status="waiting_action">
                    <tool_call call_id="aaa" name="ask_user" pending_action="respond">Some question</tool_call>
                    <tool_call call_id="bbb" name="shell" pending_action="approve">rm -rf ./directory</tool_call>
                </subtask_result>

                Followup tool call:
                    ContinueSubtask(
                        subtask_id=1,
                        message=[
                            SubtaskToolRespond(call_id="aaa", answer="Use ..."),
                            SubtaskToolApprove(call_id="bbb", decision="approved"),
                        ]
                    )

            When subtask raise an error:
                <subtask_result subtask_id="1" status="error">
                    error messages
                </subtask_result>
        """
        async def resolve_continue_subtask(task: AgentTask, continue_message: str | list[SubtaskToolRespond | SubtaskToolApprove]):
            if isinstance(continue_message, str):
                task.tool_calls.discard_pendings()
                task.messages.append(UserMessage(content=continue_message))
                return
            for message in continue_message:
                match message:
                    case SubtaskToolRespond(call_id=call_id, answer=answer):
                        task.tool_calls.apply_user_response(call_id, answer)
                    case SubtaskToolApprove(call_id=call_id, decision=decision):
                        task.tool_calls.approve(call_id, decision == "approved")

        def compose_subtask_result(subtask: AgentTask,
                                   task_result: TaskError | TaskInterrupted | TaskWaitingAction | TaskFinished) -> str:
            root = ET.Element("subtask_result", {"subtask_id": str(subtask.id)})
            match task_result:
                case TaskFinished(summary=summary):
                    root.attrib["status"] = "finished"
                    root.text = AnyXml.RawText(summary)
                case TaskWaitingAction(messages=messages):
                    root.attrib["status"] = "waiting_action"
                    for message in messages:
                        assert is_agent_tool_metadata(message.metadata)
                        tool_call_elem = ET.SubElement(root, "tool_call", {
                            "call_id": message.call_id,
                            "name": message.name,
                        })
                        if pending_action := message.metadata.get("pending_action"):
                            tool_call_elem.attrib["pending_action"] = pending_action
                        tool_call_elem.text = AnyXml.RawText(json.dumps(message.arguments, ensure_ascii=False))
                case TaskError(event=event):
                    root.attrib["status"] = "error"
                    root.text = AnyXml.RawText(event.error)
                case TaskInterrupted():
                    root.attrib["status"] = "interrupted"

            if todos := subtask.todos:
                progress_elem = ET.SubElement(root, "progress")
                for todo in todos:
                    step_elem = ET.SubElement(progress_elem, "step", {"status": todo.status})
                    step_elem.text = AnyXml.RawText(todo.description)

            return AnyXml.tostring(root)

        async with db_context() as db_session:
            subtask_service = SubtaskService(db_session)
            match action:
                case NewSubtask(instruction=instruction, agent_id=agent_id):
                    subtask = await subtask_service.create_subtask(subtask_schemas.SubtaskCreate(
                        instruction=instruction,
                        task_id=self._ctx.task_id,
                        agent_id=agent_id,
                    ))
                    task = await create_agent_task_from_subtask(subtask)
                case ContinueSubtask(subtask_id=subtask_id, agent_id=agent_id, message=message):
                    subtask = await subtask_service.get_subtask_by_id(subtask_id)
                    if agent_id is not None: subtask.agent_id = agent_id
                    task = await create_agent_task_from_subtask(subtask)
                    await resolve_continue_subtask(task, message)

        if subtask.agent_id is None:
            raise ValueError("The agent_id of this subtask is null, please pass a new agent_id when continuing this subtask.")

        try:
            await task.persist()
            task_result = await task.run_until_done()
            return compose_subtask_result(task, task_result)
        except asyncio.CancelledError:
            await task.stop()
            raise
        finally:
            await asyncio.shield(task.persist())
