"""
Context Loader — Dynamic injection and revocation of tools and artifacts.

Principle: Least privilege applied to tools. An agent receives only what
it needs for a specific task, uses it, and gives it back.

Usage:
    from abi_core.common.context_loader import load_agent_context, unload_agent_context

    # Inject tools and artifacts
    ctx = await load_agent_context(
        mcp_tool_names=["query_sales_db"],
        library_tool_names=["ShellTool"],
        artifact_keys=["task_1/data.json"],
    )
    agent.tools.extend(ctx["mcp_tools"] + ctx["lib_tools"])

    # After task — revoke
    await unload_agent_context(
        agent=agent,
        tool_names=["query_sales_db", "ShellTool"],
        cleanup_artifacts=True,
        workspace="/app/workspace",
    )
"""

import os
import shutil
from typing import Any, Dict, List, Optional

from abi_core.common.utils import abi_logging


async def load_agent_context(
    mcp_tool_names: Optional[List[str]] = None,
    library_tool_names: Optional[List[str]] = None,
    artifact_keys: Optional[List[str]] = None,
    workspace: str = "/app/workspace",
    ephemeral: bool = False,
) -> Dict[str, Any]:
    """Load dynamic context for any agent.

    Resolves MCP tools, library tools, and downloads artifacts.
    Each agent decides how to use the returned context.

    Args:
        mcp_tool_names: MCP tool names to resolve via MCPToolkit.
        library_tool_names: Library tool names to import dynamically.
        artifact_keys: MinIO keys to download to workspace.
        workspace: Local directory for artifacts.
        ephemeral: Mark tools as one-time use (metadata flag).

    Returns:
        {
            "mcp_tools": [LangChain tools from MCP],
            "lib_tools": [LangChain tools from libraries],
            "artifacts": [local file paths],
            "ephemeral": bool,
        }
    """
    result = {
        "mcp_tools": [],
        "lib_tools": [],
        "artifacts": [],
        "ephemeral": ephemeral,
    }

    os.makedirs(workspace, exist_ok=True)

    # ── MCP Tools ──────────────────────────────────────────────
    if mcp_tool_names:
        try:
            from langchain_core.tools import StructuredTool
            from abi_core.common.semantic_tools import MCPToolkit

            toolkit = MCPToolkit()
            for name in mcp_tool_names:
                async def _call_mcp(tool_name=name, **kwargs):
                    return await toolkit.call(tool_name, **kwargs)

                result["mcp_tools"].append(StructuredTool.from_function(
                    coroutine=_call_mcp,
                    name=name,
                    description=f"MCP tool: {name}",
                ))
            abi_logging(f"[🔧] Loaded {len(result['mcp_tools'])} MCP tools: {mcp_tool_names}")
        except Exception as e:
            abi_logging(f"[⚠️] Failed to load MCP tools: {e}")

    # ── Library Tools ──────────────────────────────────────────
    if library_tool_names:
        try:
            from abi_core.common.library_tools import zombie_tools, BASE_TOOLS

            # Map tool names to base tools
            base_by_name = {t.name: t for t in BASE_TOOLS}
            for name in library_tool_names:
                if name in base_by_name:
                    result["lib_tools"].append(base_by_name[name])
                else:
                    abi_logging(f"[⚠️] Library tool '{name}' not found")
            abi_logging(f"[🔧] Loaded {len(result['lib_tools'])} library tools")
        except Exception as e:
            abi_logging(f"[⚠️] Failed to load library tools: {e}")

    # ── Artifacts ──────────────────────────────────────────────
    if artifact_keys:
        try:
            from abi_core.common.artifact_store import download_artifacts
            result["artifacts"] = await download_artifacts(
                keys=artifact_keys,
                workspace=workspace,
            )
        except Exception as e:
            abi_logging(f"[⚠️] Failed to download artifacts: {e}")

    return result


async def unload_agent_context(
    agent: Any = None,
    tool_names: Optional[List[str]] = None,
    cleanup_artifacts: bool = False,
    workspace: str = "/app/workspace",
) -> None:
    """Revoke tools and clean up artifacts from an agent.

    Args:
        agent: AbiAgent instance to remove tools from.
        tool_names: Names of tools to revoke.
        cleanup_artifacts: If True, delete all files in workspace.
        workspace: Directory to clean up.
    """
    # ── Revoke tools ───────────────────────────────────────────
    if agent and tool_names and hasattr(agent, "tools"):
        names_set = set(tool_names)
        before = len(agent.tools)
        agent.tools = [t for t in agent.tools if t.name not in names_set]
        removed = before - len(agent.tools)
        if removed:
            abi_logging(f"[🔒] Revoked {removed} tools from {agent.agent_name}: {tool_names}")

    # ── Clean up artifacts ─────────────────────────────────────
    if cleanup_artifacts and os.path.exists(workspace):
        try:
            for f in os.listdir(workspace):
                path = os.path.join(workspace, f)
                if os.path.isfile(path):
                    os.remove(path)
                elif os.path.isdir(path):
                    shutil.rmtree(path)
            abi_logging(f"[🧹] Cleaned workspace: {workspace}")
        except Exception as e:
            abi_logging(f"[⚠️] Failed to clean workspace: {e}")


def build_execution_prompt(
    query: str,
    tools: Optional[List] = None,
    artifacts: Optional[List[str]] = None,
    workspace: str = "/app/workspace",
) -> str:
    """Build a dynamic execution prompt from query, tools, and artifacts.

    Generates a clean prompt that describes the task, lists available tools
    with their descriptions, and includes artifact paths if present.

    Args:
        query: The task description.
        tools: LangChain tools (each with .name and .description).
        artifacts: Local file paths of downloaded artifacts.
        workspace: Workspace directory path.

    Returns:
        Formatted prompt string ready for LLM.
    """
    parts = [f"Task: {query}"]

    if tools:
        tools_desc = "\n".join(f"- {t.name}: {t.description}" for t in tools)
        parts.append(f"\nAvailable tools:\n{tools_desc}")

    if artifacts:
        arts = "\n".join(f"  - {p}" for p in artifacts)
        parts.append(f"\nArtifacts in workspace ({workspace}):\n{arts}")

    return "\n".join(parts)
