import logging
import os
from typing import Any, Optional, List

from pydantic_ai import Agent, RunContext
from .agent_utilities import (
    read_md_file,
    append_to_md_file,
    register_a2a_peer as register_a2a_peer_util,
    create_new_skill,
    load_mcp_config,
    save_mcp_config,
    delete_skill_from_disk,
    write_skill_md,
    search_memory as search_memory_util,
    delete_memory_entry as delete_memory_entry_util,
    compress_memory as compress_memory_util,
    delete_a2a_peer as delete_a2a_peer_util,
    list_a2a_peers as list_a2a_peers_util,
    create_memory as create_memory_util,
    list_scheduled_tasks as list_scheduled_tasks_util,
    delete_scheduled_task as delete_scheduled_task_util,
    schedule_task as schedule_task_util,
    read_skill_md,
)
from .base_utilities import to_boolean

DEFAULT_WORKSPACE_TOOLS = to_boolean(string=os.environ.get("WORKSPACE_TOOLS", "True"))
DEFAULT_A2A_TOOLS = to_boolean(string=os.environ.get("A2A_TOOLS", "True"))
DEFAULT_SCHEDULER_TOOLS = to_boolean(string=os.environ.get("SCHEDULER_TOOLS", "True"))
DEFAULT_DYNAMIC_TOOLS = to_boolean(string=os.environ.get("DYNAMIC_TOOLS", "True"))

logger = logging.getLogger(__name__)

__version__ = "0.2.31"


def register_agent_tools(agent: Agent):
    """Register universal workspace, A2A, and scheduler tools to the agent."""

    # System prompt is now handled purely statically at initialization inside create_agent.

    # --- Workspace Tools ---
    if DEFAULT_WORKSPACE_TOOLS:

        @agent.tool
        async def read_workspace_file(ctx: RunContext[Any], filename: str) -> str:
            """Read content of any .md file in workspace (IDENTITY.md, CRON.md, etc.).
            Read your system prompt, identity, memory, scheduled cron jobs, cron log, and mcp_config.json
            """
            return read_md_file(filename)

        @agent.tool
        async def append_note_to_file(
            ctx: RunContext[Any], filename: str, text: str
        ) -> str:
            """Append a short note or section to a workspace .md file."""
            append_to_md_file(filename, text)
            return f"Appended to {filename}"

        @agent.tool
        async def create_memory(ctx: RunContext[Any], text: str) -> str:
            """
            Save important decisions, outcomes, user preferences, critical
            information, or information the user explicitly requests to long-term memory (MEMORY.md).
            """
            return create_memory_util(text)

        @agent.tool
        async def search_memory(ctx: RunContext[Any], query: str) -> str:
            """Search through long-term memory (MEMORY.md) for specific information."""
            return search_memory_util(query)

        @agent.tool
        async def delete_memory_entry(ctx: RunContext[Any], index: int) -> str:
            """Delete a specific memory entry by line index (1-based). Use search_memory first to find the index."""
            return delete_memory_entry_util(index)

        @agent.tool
        async def compress_memory(ctx: RunContext[Any], max_entries: int = 50) -> str:
            """Compress or prune long-term memory to keep it concise and relevant."""
            return compress_memory_util(max_entries)

    # --- A2A Peer Tools ---
    if DEFAULT_A2A_TOOLS:

        @agent.tool
        async def list_a2a_peers(ctx: RunContext[Any]) -> str:
            """List all known A2A peer agents."""
            return list_a2a_peers_util()

        @agent.tool
        async def register_a2a_peer(
            ctx: RunContext[Any],
            name: str,
            url: str,
            description: str = "",
            capabilities: str = "",
            auth: str = "none",
        ) -> str:
            """Register or update another A2A agent this agent can call."""
            return register_a2a_peer_util(name, url, description, capabilities, auth)

        @agent.tool
        async def delete_a2a_peer(ctx: RunContext[Any], name: str) -> str:
            """Remove an A2A peer agent from the local registry."""
            return delete_a2a_peer_util(name)

    # --- Scheduler Tools ---
    if DEFAULT_SCHEDULER_TOOLS:

        @agent.tool
        async def schedule_task(
            ctx: RunContext[Any],
            task_id: str,
            name: str,
            interval_minutes: int,
            prompt: str,
        ) -> str:
            """Schedule a task to run periodically (persists in CRON.md)."""
            return schedule_task_util(task_id, name, interval_minutes, prompt)

        @agent.tool
        async def list_tasks(ctx: RunContext[Any]) -> str:
            """List all active periodic tasks."""
            return list_scheduled_tasks_util()

        @agent.tool
        async def delete_task(ctx: RunContext[Any], task_id: str) -> str:
            """Permanently remove a scheduled task by ID."""
            return delete_scheduled_task_util(task_id)

    # --- Dynamic Extension Tools ---
    if DEFAULT_DYNAMIC_TOOLS:

        @agent.tool
        async def update_mcp_config(
            ctx: RunContext[Any],
            server_name: str,
            server_url: str,
            transport: str = "sse",
        ) -> str:
            """Add or update an MCP server entry in mcp_config.json (auto-loaded on next run)."""
            config = load_mcp_config()
            config.setdefault("mcpServers", {})[server_name] = {
                "url": server_url,
                "transport": transport,
            }
            save_mcp_config(config)
            return f"✅ MCP server '{server_name}' added/updated. It will be available on next run."

        @agent.tool
        async def create_skill(
            ctx: RunContext[Any],
            name: str,
            description: str,
            when_to_use: str = "",
            how_to_use: str = "",
        ) -> str:
            """Create a brand-new skill folder + SKILL.md that will be auto-loaded on next run."""
            return create_new_skill(name, description, when_to_use, how_to_use)

        @agent.tool
        async def delete_skill(ctx: RunContext[Any], name: str) -> str:
            """Delete a skill folder from the workspace. Only works for workspace skills."""
            return delete_skill_from_disk(name)

        @agent.tool
        async def edit_skill(ctx: RunContext[Any], name: str, new_content: str) -> str:
            """
            Overwrite the SKILL.md of an existing workspace skill.
            Use this to refine a skill's logic, description, or examples.
            """
            return write_skill_md(name, new_content)

        @agent.tool
        async def get_skill_content(ctx: RunContext[Any], name: str) -> str:
            """Read the current SKILL.md of a workspace skill to prepare for editing."""
            return read_skill_md(name)

        @agent.tool
        async def get_mcp_reference(ctx: RunContext[Any], agent_template: str) -> str:
            """
            Retrieve the documentation for an MCP agent template.
            Use this to discover available 'enabled_tools' tags and configuration details.

            Args:
                agent_template: The name of the agent template. Options: ["adguard-home", "ansible-tower", "archivebox", "arr", "audio-transcriber", "container-manager", "documentdb", "github", "gitlab", "jellyfin", "leanix", "mealie", "media-downloader", "microsoft", "nextcloud", "plane", "portainer", "repository-manager", "searxng", "servicenow", "stirlingpdf", "systems-manager", "tunnel-manager", "vector", "wger"]
            """
            from universal_skills.skill_utilities import resolve_mcp_reference
            from pathlib import Path

            # Try direct name, then with -mcp suffix, then with -agent suffix
            # resolve_mcp_reference handles extension mapping logic
            queries = [
                f"{agent_template}.md",
                f"{agent_template}-api.md",
                f"{agent_template}-mcp.md",
                f"{agent_template}-agent.md",
            ]

            md_path = None
            for query in queries:
                md_path = resolve_mcp_reference(query)
                if md_path:
                    break

            if not md_path:
                return f"Error: Reference documentation for '{agent_template}' not found. Tried: {', '.join(queries)}"

            path = Path(md_path)
            if not path.exists():
                return f"Error: Documentation file {md_path} does not exist."

            return path.read_text(encoding="utf-8")

        @agent.tool
        async def spawn_agent(
            ctx: RunContext[Any],
            agent_template: Optional[str] = None,
            prompt: str = "Please assist.",
            mcp_config: Optional[str] = None,
            mcp_url: Optional[str] = None,
            enabled_tools: Optional[List[str]] = None,
            custom_skills_directory: Optional[str] = None,
            name: Optional[str] = None,
            system_prompt: Optional[str] = None,
            provider: Optional[str] = None,
            model_id: Optional[str] = None,
            base_url: Optional[str] = None,
            api_key: Optional[str] = None,
        ) -> str:
            """
            Spawn a new Pydantic AI sub-agent dynamically to handle a specific task.
            Use this when you need to delegate a task to an agent with a specialized external configuration.

            Args:
                agent_template: Pre-defined agent configuration from the mcp-client skill. Options: ["adguard-home", "ansible-tower", "archivebox", "arr", "audio-transcriber", "container-manager", "documentdb", "github", "gitlab", "jellyfin", "leanix", "mealie", "media-downloader", "microsoft", "nextcloud", "plane", "portainer", "repository-manager", "searxng", "servicenow", "stirlingpdf", "systems-manager", "tunnel-manager", "vector", "wger"]
                prompt: The instruction to send to the spawned agent.
                mcp_config: Path to a specific MCP config JSON file or a filename from the mcp-client skill's references.
                enabled_tools: Optional list of tool tags to enable (e.g., ["BRANCHESTOOL", "COMMITSTOOL"]). All others will be disabled if the config supports filtering. Use 'get_mcp_reference' to discover available tags for a template.
                name: Override the agent's name.
                system_prompt: Override the agent's system prompt.
            """
            from .agent_utilities import (
                create_agent,
                get_mcp_config_path,
                get_system_prompt_from_reference,
            )
            from universal_skills.skill_utilities import resolve_mcp_reference
            import os

            AGENT_TEMPLATES = {
                "adguard-home": {
                    "name": "AdGuard Home Agent",
                    "mcp_config": "adguard-home-agent.json",
                },
                "ansible-tower": {
                    "name": "Ansible Tower Agent",
                    "mcp_config": "ansible-tower-mcp.json",
                },
                "archivebox": {
                    "name": "ArchiveBox Agent",
                    "mcp_config": "archivebox-api.json",
                },
                "arr": {
                    "name": "Arr Agent",
                    "mcp_config": "arr-mcp.json",
                },
                "audio-transcriber": {
                    "name": "Audio Transcriber Agent",
                    "mcp_config": "audio-transcriber-mcp.json",
                },
                "container-manager": {
                    "name": "Container Manager Agent",
                    "mcp_config": "container-manager-mcp.json",
                },
                "documentdb": {
                    "name": "DocumentDB Agent",
                    "mcp_config": "documentdb-mcp.json",
                },
                "github": {
                    "name": "GitHub Agent",
                    "mcp_config": "github-mcp.json",
                },
                "gitlab": {
                    "name": "GitLab Agent",
                    "mcp_config": "gitlab-api.json",
                },
                "jellyfin": {
                    "name": "Jellyfin Agent",
                    "mcp_config": "jellyfin-mcp.json",
                },
                "leanix": {
                    "name": "LeanIX Agent",
                    "mcp_config": "leanix-agent.json",
                },
                "mealie": {
                    "name": "Mealie Agent",
                    "mcp_config": "mealie-mcp.json",
                },
                "media-downloader": {
                    "name": "Media Downloader Agent",
                    "mcp_config": "media-downloader-mcp.json",
                },
                "microsoft": {
                    "name": "Microsoft 365 Agent",
                    "mcp_config": "microsoft-agent.json",
                },
                "nextcloud": {
                    "name": "Nextcloud Agent",
                    "mcp_config": "nextcloud-agent.json",
                },
                "plane": {
                    "name": "Plane Agent",
                    "mcp_config": "plane-mcp.json",
                },
                "portainer": {
                    "name": "Portainer Agent",
                    "mcp_config": "portainer-agent.json",
                },
                "repository-manager": {
                    "name": "Repository Manager Agent",
                    "mcp_config": "repository-manager.json",
                },
                "searxng": {
                    "name": "SearXNG Agent",
                    "mcp_config": "searxng-mcp.json",
                },
                "servicenow": {
                    "name": "ServiceNow Agent",
                    "mcp_config": "servicenow-api.json",
                },
                "stirlingpdf": {
                    "name": "Stirling PDF Agent",
                    "mcp_config": "stirlingpdf-agent.json",
                },
                "systems-manager": {
                    "name": "Systems Manager Agent",
                    "mcp_config": "systems-manager-mcp.json",
                },
                "tunnel-manager": {
                    "name": "Tunnel Manager Agent",
                    "mcp_config": "tunnel-manager-mcp.json",
                },
                "vector": {
                    "name": "Vector Agent",
                    "mcp_config": "vector-mcp.json",
                },
                "wger": {
                    "name": "Wger Fitness Agent",
                    "mcp_config": "wger-agent.json",
                },
            }

            template_config = (
                AGENT_TEMPLATES.get(agent_template, {}) if agent_template else {}
            )

            # Priority: 1. Explicit arguments 2. Template Database (Dynamic Prompt) 3. OS Fallbacks
            final_name = name or template_config.get("name", "SpawnedAgent")

            # Dynamic prompt lookup for templates
            dynamic_prompt = (
                get_system_prompt_from_reference(agent_template)
                if agent_template
                else None
            )
            final_system_prompt = (
                system_prompt
                or dynamic_prompt
                or template_config.get("system_prompt", "You are a helpful assistant.")
            )

            # Resolve MCP configuration
            final_mcp_config = mcp_config or template_config.get("mcp_config")
            if final_mcp_config:
                resolved_path = resolve_mcp_reference(final_mcp_config)
                if resolved_path:
                    final_mcp_config = resolved_path
                else:
                    final_mcp_config = get_mcp_config_path(final_mcp_config)

            # Tailored tool loading: Temporarily inject tools to enable into os.environ
            # This works because load_mcp_servers expands env vars from the current environment.
            original_env = {}
            if enabled_tools:
                for tool in enabled_tools:
                    original_env[tool] = os.environ.get(tool)
                    os.environ[tool] = "True"

                # We assume if enabled_tools is provided, we might want to disable others
                # that are typically defaulted to True in the JSON configs but support env override.
                # However, since we don't know all possible tools, we rely on the user to
                # specify what they want. For GitLab specifically, others are defaulted to True
                # in the reference JSON, but can be overridden by env vars.
                # To truly "only" enable specific ones, the JSON would need ${VAR:-True}
                # but if we set the other vars to "False", they will be overridden.

            # Create the sub-agent
            try:
                sub_agent = create_agent(
                    provider=provider
                    or template_config.get("provider")
                    or os.getenv("PROVIDER", "openai"),
                    model_id=model_id
                    or template_config.get("model_id")
                    or os.getenv("MODEL_ID", "gpt-4o"),
                    base_url=base_url
                    or template_config.get("base_url")
                    or os.getenv("LLM_BASE_URL"),
                    api_key=api_key
                    or template_config.get("api_key")
                    or os.getenv("LLM_API_KEY"),
                    mcp_config=final_mcp_config,
                    mcp_url=mcp_url
                    or template_config.get("mcp_url")
                    or os.getenv("MCP_URL"),
                    name=final_name,
                    system_prompt=final_system_prompt,
                    # elicitation_callback is now handled client-side per agent
                    debug=True,
                )
            finally:
                # Restore environment
                for tool, original_val in original_env.items():
                    if original_val is None:
                        del os.environ[tool]
                    else:
                        os.environ[tool] = original_val

            # Execute the prompt
            logger.info(
                f"Sub-agent run started: {agent_template}. Prompt: {prompt[:100]}..."
            )

            # Retrieve elicitation queue from current context deps if available
            elicitation_queue = None
            if hasattr(ctx, "deps") and isinstance(ctx.deps, dict):
                elicitation_queue = ctx.deps.get("elicitation_queue")

            try:
                # Pass elicitation_queue to sub-agent to preserve context
                result = await sub_agent.run(
                    prompt, deps={"elicitation_queue": elicitation_queue}
                )
                logger.info(f"Sub-agent run completed: {agent_template}")

                # Safe access to result data
                data = getattr(result, "data", None)
                if data is None:
                    # Fallback for some versions or result types
                    data = getattr(result, "result", result)
                return str(data)
            except Exception as e:
                import traceback

                error_msg = f"Error in sub-agent run: {e}\n{traceback.format_exc()}"
                logger.error(error_msg)
                return error_msg
