import platform
import xml.etree.ElementTree as ET
from collections import namedtuple
from dataclasses import asdict
from typing import Self
from anyxml import AnyXml
from loguru import logger
from dais_sdk.tool import Toolset
from dais_sdk.types import Message, ToolDef, ToolFn
from src.db import db_context
from src.schemas import (
    agent as agent_schemas,
    workspace as workspace_schemas,
    provider as provider_schemas,
    skill as skill_schemas,
)
from src.schemas.tasks import runtime as task_runtime_schemas
from src.services.agent import AgentService
from src.services.workspace import WorkspaceService
from src.services.llm_model import LlmModelService
from src.services.provider import ProviderService
from src.settings import use_app_setting_manager
from .persistence import create_agent_context_persistence
from .aliases import BuiltInToolAliases
from .models import AgentContextResource, AgentContextPersistence
from ..notes import NoteMaterializer
from ..tool import use_mcp_toolset_manager, BuiltinToolsetManager, BuiltinToolsetContext, McpToolsetManager
from ..tool.types import is_tool_metadata
from ..prompts import (
    BASE_INSTRUCTION,
    DEFAULT_BASE_ROLE,
    APPENDIX_TEMPLATE,
    FAILED_TO_LOAD_NOTES_INDEX,
    NO_AVAILABLE_SKILLS,
    NO_AVAILABLE_AGENTS,
    NO_AVAILABLE_TOOLSETS,
    NO_WORKSPACE_INSTRUCTION,
    NO_AGENT_INSTRUCTION,
)
from ..types import ContextUsage


class AgentContext:
    _logger = logger.bind(name="AgentContext")

    def __init__(self,
                 task_id: int,
                 task_type: task_runtime_schemas.TaskType,
                 *,
                 messages: list[Message],
                 resource: AgentContextResource,
                 usage: ContextUsage,
                 persistence: AgentContextPersistence,
                 builtin_toolset_manager: BuiltinToolsetManager,
                 mcp_toolset_manager: McpToolsetManager):
        self.task_id = task_id
        self.task_type = task_type
        self._resource = resource
        self._messages = messages

        self._usage = usage
        self._persistence = persistence
        self._builtin_toolset_manager = builtin_toolset_manager
        self._mcp_toolset_manager = mcp_toolset_manager
        self._builtin_tool_aliases = BuiltInToolAliases(builtin_toolset_manager)

    @classmethod
    async def create(cls, task: task_runtime_schemas.TaskRuntimeContext) -> Self:
        assert task.agent_id is not None

        async with db_context() as db_session:
            agent = await AgentService(db_session).get_agent_by_id(task.agent_id)
            workspace = await WorkspaceService(db_session).get_workspace_by_id(task.workspace_id)
            skills = workspace.usable_skills

            assert agent.model_id is not None
            model = await LlmModelService(db_session).get_model_by_id(agent.model_id)
            provider = await ProviderService(db_session).get_provider_by_id(model.provider_id)

        usage = task.usage
        usage.max_tokens = model.context_size
        usage = ContextUsage(**asdict(usage))
        messages = task.messages

        builtin_toolset_manager = await BuiltinToolsetManager.create(
            BuiltinToolsetContext(
                task.id,
                workspace.id,
                workspace.directory,
            )
        )
        mcp_toolset_manager = use_mcp_toolset_manager()

        persistence = create_agent_context_persistence(task)
        return cls(task.id,
                   task.type,
                   messages=messages,
                   resource=AgentContextResource(
                       workspace=workspace_schemas.WorkspaceRead.model_validate(workspace),
                       agent=agent_schemas.AgentRead.model_validate(agent),
                       provider=provider_schemas.ProviderRead.model_validate(provider),
                       model=provider_schemas.LlmModelRead.model_validate(model),
                       skills=[skill_schemas.SkillBrief.model_validate(skill) for skill in skills],
                   ),
                   usage=usage,
                   persistence=persistence,
                   builtin_toolset_manager=builtin_toolset_manager,
                   mcp_toolset_manager=mcp_toolset_manager)

    @staticmethod
    def _format_skills(skills: list[skill_schemas.SkillBrief]) -> str:
        if len(skills) == 0: return NO_AVAILABLE_SKILLS
        root = ET.Element("skills")
        for skill in skills:
            skill_elem = ET.SubElement(root, "skill", attrib={"id": str(skill.id), "name": skill.name})
            skill_elem.text = AnyXml.RawText(skill.description)
        return AnyXml.tostring(root)

    @staticmethod
    def _format_agents(agents: list[agent_schemas.AgentBrief]) -> str:
        if len(agents) == 0: return NO_AVAILABLE_AGENTS
        root = ET.Element("agents")
        for agent in agents:
            agent_elem = ET.SubElement(root, "agent", attrib={"id": str(agent.id), "name": agent.name})
            if agent.description:
                agent_elem.text = AnyXml.RawText(agent.description)
        return AnyXml.tostring(root)

    @staticmethod
    def _format_toolsets(toolsets: list[Toolset]) -> str:
        if len(toolsets) == 0: return NO_AVAILABLE_TOOLSETS
        root = ET.Element("toolsets")
        for toolset in toolsets:
            toolset_elem = ET.Element("toolset", attrib={"name": toolset.name})
            if toolset.description is not None:
                toolset_elem.text = AnyXml.RawText(toolset.description)
        return AnyXml.tostring(root)

    ResolvedInstructions = namedtuple("ResolvedInstructions", ["base", "workspace", "agent"])
    def _resolve_instructions(self) -> ResolvedInstructions:
        workspace_instruction = self._resource.workspace.instruction or NO_WORKSPACE_INSTRUCTION
        agent_instruction = self._resource.agent.instruction or NO_AGENT_INSTRUCTION

        builtin_tool_aliases = self._builtin_tool_aliases
        resolved_base_instruction = builtin_tool_aliases.substitute(BASE_INSTRUCTION)
        resolved_workspace_instruction = builtin_tool_aliases.substitute(workspace_instruction)
        resolved_agent_instruction = builtin_tool_aliases.substitute(agent_instruction)
        return AgentContext.ResolvedInstructions(resolved_base_instruction, resolved_workspace_instruction, resolved_agent_instruction)

    @property
    def usage(self) -> ContextUsage:
        return self._usage

    @property
    def skills(self) -> list[skill_schemas.SkillBrief]:
        return [skill for skill in self._resource.skills if skill.is_enabled]

    @property
    def toolsets(self) -> list[Toolset]:
        return [*self._builtin_toolset_manager.toolsets,
                *self._mcp_toolset_manager.toolsets]

    @property
    def provider(self) -> provider_schemas.ProviderRead:
        return self._resource.provider

    @property
    def model(self) -> provider_schemas.LlmModelRead:
        return self._resource.model

    @property
    def messages(self) -> list[Message]: return self._messages
    @messages.setter
    def messages(self, new_messages: list[Message]): self._messages = new_messages

    async def filter_usable_tool_ids(self) -> set[int] | None:
        from ..tool import BuiltinToolset, OrchestrationToolset, UserInteractionToolset
        def find_builtin_tool_id(toolset_cls: type[BuiltinToolset], target_tool: ToolFn) -> int | None:
            target_toolset = None
            for toolset in self._builtin_toolset_manager.toolsets:
                assert isinstance(toolset, BuiltinToolset)
                if toolset.internal_key() == toolset_cls.internal_key():
                    target_toolset = toolset
            if target_toolset is None: return None
            for tool_def in target_toolset.get_tools():
                if tool_def.executes(target_tool):
                    assert is_tool_metadata(tool_def.metadata)
                    return tool_def.metadata["id"]
            return None

        workspace_usable_tool_ids = {tool.id for tool in self._resource.workspace.usable_tools}
        agent_usable_tool_ids = {tool.id for tool in self._resource.agent.usable_tools}

        if len(workspace_usable_tool_ids) == 0 and len(agent_usable_tool_ids) == 0:
            # both workspace and agent have no usable tools configured, return None meaning no need to filter
            return None

        if self.task_type == "subtask":
            # remove "subtask" tool when the current task_type is "subtask"
            subtask_tool_id = find_builtin_tool_id(OrchestrationToolset, OrchestrationToolset.subtask)
            if subtask_tool_id is not None:
                if subtask_tool_id in workspace_usable_tool_ids: workspace_usable_tool_ids.remove(subtask_tool_id)
                if subtask_tool_id in agent_usable_tool_ids: agent_usable_tool_ids.remove(subtask_tool_id)

        if self.task_type == "schedule":
            # remove "show_plan" tool when the current task_type is "schedule"
            show_plan_tool_id = find_builtin_tool_id(UserInteractionToolset, UserInteractionToolset.show_plan)
            if show_plan_tool_id is not None:
                if show_plan_tool_id in workspace_usable_tool_ids: workspace_usable_tool_ids.remove(show_plan_tool_id)
                if show_plan_tool_id in agent_usable_tool_ids: agent_usable_tool_ids.remove(show_plan_tool_id)

        if len(workspace_usable_tool_ids) == 0:
            return agent_usable_tool_ids
        if len(agent_usable_tool_ids) == 0:
            return workspace_usable_tool_ids
        return workspace_usable_tool_ids & agent_usable_tool_ids

    async def compose_system_instruction(self) -> str:
        settings = use_app_setting_manager().settings
        available_skills = AgentContext._format_skills(self.skills)
        available_toolsets = AgentContext._format_toolsets(self.toolsets)
        resolved_instructions = self._resolve_instructions()

        resolved_agents_appendix = APPENDIX_TEMPLATE.format(
            id="D",
            title="Available Agents",
            content=AgentContext._format_agents(self._resource.workspace.usable_agents)
        ) if self.task_type != "subtask" else "" # does not inject agents info under "subtask"

        notes_index = await NoteMaterializer.get_notes_index(self._resource.workspace.id)
        resolved_notes_index = notes_index if notes_index is not None else FAILED_TO_LOAD_NOTES_INDEX

        return resolved_instructions.base.format(
            base_role=DEFAULT_BASE_ROLE,
            os_platform=platform.system(),
            user_language=settings.reply_language,
            available_skills=available_skills,
            available_toolsets=available_toolsets,
            runtime_appendices=resolved_agents_appendix,
            workspace_notes_index=resolved_notes_index,
            workspace_name=self._resource.workspace.name,
            workspace_directory=self._resource.workspace.directory,
            workspace_instruction=resolved_instructions.workspace,
            agent_role=self._resource.agent.name,
            agent_instruction=resolved_instructions.agent,
        ).strip()

    def find_tool(self, tool_name: str) -> ToolDef | None:
        for toolset in self.toolsets:
            for tool in toolset.get_tools():
                if tool.name == tool_name:
                    return tool
        return None

    async def persist(self) -> task_runtime_schemas.TaskRuntimeContext:
        return await self._persistence.persist(
            self.task_id,
            self._messages,
            self._usage,
        )
