"""Role-based tool filtering and progressive tool loading.

Filters the tool set presented to the LLM based on the agent's role
and task type. Reduces token overhead and improves selection accuracy.

Implements 3-tier progressive loading:
  Tier 1: One-line summaries (~100 tokens) — loaded in system prompt
  Tier 2: Full documentation (<5k tokens) — loaded on first tool use
  Tier 3: Resources (examples, schemas) — on-demand
"""

from __future__ import annotations


# ---------------------------------------------------------------------------
# Tier-1: Compact one-line tool summaries for system prompt injection.
# Each description should be <=60 chars for minimal token overhead.
# ---------------------------------------------------------------------------
TOOL_SUMMARIES: dict[str, str] = {
    "web_search": "Search the web using keywords",
    "web_fetch": "Fetch and extract text from a URL",
    "read_file": "Read contents of a workspace file",
    "write_file": "Write content to a workspace file",
    "edit_file": "Edit sections of an existing file",
    "list_dir": "List files and directories in workspace",
    "exec": "Execute a shell command and return output",
    "message": "Send a message or deliverable to control plane",
    "browser_navigate": "Navigate the browser to a URL",
    "browser_snapshot": "Capture current page state (DOM + screenshot)",
    "browser_click": "Click an element on the page",
    "browser_type": "Type text into a browser input field",
    "browser_wait": "Wait for a browser element or condition",
    "browser_back": "Navigate browser back one page",
    "browser_screenshot": "Take a screenshot of the browser",
    "browser_execute": "Execute JavaScript in the browser",
    "browser_select": "Select an option from a dropdown",
    "browser_checkbox": "Toggle a checkbox element",
    "browser_press": "Press a keyboard key in the browser",
    "browser_fill_form": "Fill multiple form fields at once",
    "browser_upload_file": "Upload a file via browser file input",
    "browser_scroll": "Scroll the browser page",
    "browser_hover": "Hover over a browser element",
    "browser_drag": "Drag an element to a target",
    "browser_new_tab": "Open a new browser tab",
    "browser_switch_tab": "Switch to a different browser tab",
    "browser_close_tab": "Close a browser tab",
    "download": "Download a file from a URL to workspace",
    "pdf_read": "Extract text from a PDF file",
    "transcribe": "Transcribe audio/video to text",
    "extract_tables": "Extract HTML tables from browser page",
    "extract_links": "Extract links from browser page",
    "extract_forms": "Extract form fields from browser page",
    "extract_text": "Extract visible text from browser page",
    "extract_meta": "Extract page metadata (title, desc, etc.)",
    "screenshot_analyze": "Analyze a screenshot using vision LLM",
    "image_analyze": "Analyze an image file using vision LLM",
    "python_exec": "Run Python code in a sandboxed subprocess",
    "http_request": "Make HTTP requests to REST APIs",
    "web_crawl": "Crawl multiple pages following links",
    "spreadsheet": "Create, read, modify CSV and XLSX files",
    "chart": "Generate charts as PNG images",
    "diff": "Compare two texts or files for differences",
    "archive": "Create and extract ZIP/TAR archives",
    "email": "Send and receive emails via SMTP/IMAP",
    "ask_human": "Ask the user a question and wait for reply",
    "request_approval": "Request approval before irreversible actions",
    "notify": "Send notifications via control plane or webhook",
    "page_to_pdf": "Save current browser page as PDF",
    "dom_query": "Query DOM elements by CSS selector",
    "page_search": "Grep-like text search within page",
    "sqlite": "Create and query SQLite databases",
    "template_render": "Render Jinja2 templates with variables",
    "file_convert": "Convert formats: MD/HTML/image/CSV/JSON",
    "regex": "Match, extract, replace with regex patterns",
    "json_query": "Query and transform JSON with dot-paths",
    "clipboard": "Read/write system clipboard",
    "browser_forward": "Navigate browser forward one page",
    "browser_double_click": "Double-click an element on the page",
    "browser_clear": "Clear an input field",
    "cron": "Schedule a recurring task",
    "spawn": "Spawn a sub-agent process",
}


# ---------------------------------------------------------------------------
# Tier-2: Full tool documentation registry (populated during tool setup).
# ---------------------------------------------------------------------------
_TOOL_FULL_DOCS: dict[str, str] = {}


def register_tool_docs(tool_name: str, docs: str) -> None:
    """Register full documentation for a tool (called during tool setup)."""
    _TOOL_FULL_DOCS[tool_name] = docs


def get_tool_full_docs(tool_name: str) -> str | None:
    """Return full documentation for a specific tool (tier-2).

    Called on-demand when an agent first invokes a tool, to provide
    complete parameter descriptions and usage examples.
    Returns None if tool not found.
    """
    return _TOOL_FULL_DOCS.get(tool_name)


# Role -> tool filtering profile
# "always": tools always included for this role
# "available": tools included but not prioritized
# "exclude": tools hidden from this role
ROLE_TOOL_PROFILES: dict[str, dict[str, list[str] | None]] = {
    "researcher": {
        "always": [
            "web_search", "web_fetch", "read_file", "write_file", "message",
            "download", "pdf_read", "python_exec", "extract_text", "extract_links",
            "http_request", "web_crawl", "diff", "ask_human",
            "json_query", "regex",
        ],
        "available": [
            "browser_navigate", "browser_snapshot", "browser_click",
            "browser_type", "browser_wait", "browser_back",
            "browser_scroll", "extract_tables", "extract_meta",
            "screenshot_analyze", "image_analyze", "transcribe",
            "exec", "list_dir", "edit_file",
            "spreadsheet", "chart", "archive", "notify", "request_approval",
            "dom_query", "page_search", "page_to_pdf",
            "sqlite", "template_render", "file_convert", "clipboard",
        ],
        "exclude": ["cron", "spawn"],
    },
    "writer": {
        "always": [
            "read_file", "write_file", "edit_file", "message", "list_dir", "python_exec",
            "diff", "ask_human", "template_render", "regex",
        ],
        "available": [
            "web_fetch", "web_search", "exec", "pdf_read", "image_analyze", "download",
            "spreadsheet", "chart", "archive", "request_approval",
            "json_query", "file_convert", "clipboard",
        ],
        "exclude": [
            "browser_navigate", "browser_snapshot", "browser_click",
            "browser_type", "browser_wait", "browser_back",
            "browser_screenshot", "browser_execute",
            "cron", "spawn",
        ],
    },
    "analyst": {
        "always": [
            "web_search", "web_fetch", "read_file", "exec", "message", "write_file",
            "python_exec", "download", "pdf_read", "extract_tables",
            "http_request", "spreadsheet", "chart", "diff", "ask_human",
            "sqlite", "json_query", "regex",
        ],
        "available": [
            "browser_navigate", "browser_snapshot", "list_dir", "edit_file",
            "extract_text", "extract_links", "extract_meta",
            "image_analyze", "transcribe",
            "web_crawl", "archive", "notify", "request_approval",
            "file_convert", "template_render", "clipboard",
        ],
        "exclude": ["cron", "spawn"],
    },
    "browser_operator": {
        "always": [
            "browser_navigate", "browser_snapshot", "browser_click",
            "browser_type", "browser_wait", "browser_back",
            "browser_screenshot", "browser_execute",
            "browser_select", "browser_checkbox", "browser_press",
            "browser_fill_form", "browser_upload_file",
            "browser_scroll", "browser_hover",
            "browser_new_tab", "browser_switch_tab", "browser_close_tab",
            "browser_forward", "browser_double_click", "browser_clear",
            "extract_forms", "extract_tables", "extract_text", "extract_links",
            "screenshot_analyze", "download",
            "message", "web_search", "ask_human", "request_approval",
            "dom_query", "page_search", "page_to_pdf",
        ],
        "available": [
            "read_file", "write_file", "exec", "web_fetch", "list_dir",
            "pdf_read", "python_exec", "image_analyze", "transcribe",
            "browser_drag", "extract_meta",
            "http_request", "spreadsheet", "notify",
            "clipboard", "json_query",
        ],
        "exclude": ["cron"],
    },
    "transcriber": {
        "always": [
            "transcribe", "download", "read_file", "write_file", "message",
            "python_exec", "ask_human",
        ],
        "available": [
            "web_fetch", "web_search", "exec", "pdf_read", "list_dir",
            "spreadsheet", "notify",
        ],
        "exclude": [
            "browser_navigate", "browser_snapshot", "browser_click",
            "browser_type", "cron", "spawn",
        ],
    },
    "data_extractor": {
        "always": [
            "web_search", "web_fetch", "download", "python_exec",
            "extract_tables", "extract_text", "extract_links", "extract_forms",
            "browser_navigate", "browser_snapshot", "browser_scroll",
            "read_file", "write_file", "message", "pdf_read",
            "http_request", "web_crawl", "spreadsheet",
            "dom_query", "page_search", "json_query",
        ],
        "available": [
            "browser_click", "browser_type", "browser_wait", "browser_back",
            "exec", "list_dir", "edit_file", "image_analyze", "extract_meta",
            "chart", "archive", "diff", "notify", "ask_human", "request_approval",
            "regex", "sqlite", "file_convert", "page_to_pdf",
        ],
        "exclude": ["cron", "spawn"],
    },
    "assistant": {
        "always": [
            "web_search", "web_fetch", "read_file", "write_file", "edit_file",
            "message", "exec", "download", "python_exec",
            "http_request", "spreadsheet", "email", "notify", "ask_human", "request_approval",
            "diff", "archive", "chart",
            "sqlite", "template_render", "file_convert", "regex", "json_query", "clipboard",
        ],
        "available": [
            "browser_navigate", "browser_snapshot", "browser_click",
            "browser_type", "browser_fill_form", "browser_select",
            "browser_scroll", "browser_upload_file",
            "pdf_read", "image_analyze", "transcribe",
            "web_crawl", "extract_tables", "extract_text", "extract_links",
            "list_dir",
            "dom_query", "page_search", "page_to_pdf",
            "browser_forward", "browser_double_click", "browser_clear",
        ],
        "exclude": ["cron", "spawn"],
    },
    "default": {
        "always": None,  # All tools
        "available": None,
        "exclude": [],
    },
}

# Task-type overrides that add tools regardless of role
TASK_TYPE_TOOL_ADDITIONS: dict[str, list[str]] = {
    "investigation": [
        "web_search", "web_fetch", "browser_navigate", "browser_snapshot",
        "download", "pdf_read", "extract_text", "extract_links",
        "http_request", "web_crawl", "diff",
    ],
    "coding": ["exec", "read_file", "write_file", "edit_file", "list_dir", "python_exec", "diff"],
    "browsing": [
        "browser_navigate", "browser_snapshot", "browser_click",
        "browser_type", "browser_wait", "browser_select", "browser_checkbox",
        "browser_fill_form", "browser_upload_file", "browser_scroll",
        "browser_press", "browser_hover",
        "browser_forward", "browser_double_click", "browser_clear",
        "extract_forms", "extract_tables", "screenshot_analyze",
        "ask_human", "request_approval", "dom_query", "page_search", "page_to_pdf",
    ],
    "scraping": [
        "browser_navigate", "browser_snapshot", "browser_scroll",
        "extract_tables", "extract_links", "extract_text", "extract_meta",
        "download", "python_exec", "web_fetch",
        "web_crawl", "http_request", "spreadsheet",
    ],
    "transcription": ["transcribe", "download", "python_exec", "spreadsheet"],
    "document_processing": ["pdf_read", "download", "image_analyze", "python_exec", "spreadsheet", "diff"],
    "data_analysis": [
        "python_exec", "spreadsheet", "chart", "pdf_read",
        "http_request", "diff", "read_file", "write_file",
        "sqlite", "json_query", "regex", "file_convert",
    ],
    "communication": ["email", "notify", "ask_human", "request_approval", "message"],
    "file_management": ["archive", "download", "read_file", "write_file", "list_dir", "diff", "file_convert", "template_render"],
    "monitoring": ["http_request", "web_crawl", "notify", "cron", "spreadsheet"],
}


def get_tools_for_context(
    role_tag: str | None = None,
    task_type: str | None = None,
    all_tool_names: list[str] | None = None,
) -> set[str] | None:
    """Return the set of tool names to include, or None for all tools.

    Args:
        role_tag: Agent role (e.g., "researcher", "writer", "analyst")
        task_type: Task classification (e.g., "investigation", "coding")
        all_tool_names: Full list of available tool names (for validation)

    Returns:
        Set of tool names to include, or None to include all tools.
    """
    profile = ROLE_TOOL_PROFILES.get(role_tag or "default", ROLE_TOOL_PROFILES["default"])

    # Default profile = all tools
    if profile["always"] is None:
        return None

    tools = set(profile["always"] or [])
    tools.update(profile.get("available") or [])
    tools -= set(profile.get("exclude") or [])

    # Task-type additions override exclusions
    if task_type:
        additions = TASK_TYPE_TOOL_ADDITIONS.get(task_type, [])
        tools.update(additions)

    # Always include message tool (required for deliverables)
    tools.add("message")

    # Validate against actual available tools if provided
    if all_tool_names is not None:
        available = set(all_tool_names)
        tools = tools & available

    return tools


def get_tier1_tool_descriptions(
    role_tag: str | None = None,
    task_type: str | None = None,
    all_tool_names: list[str] | None = None,
) -> str:
    """Return compact tool summaries (tier-1) for the agent's context.

    Instead of full tool documentation, returns one-line descriptions.
    Full docs are loaded on-demand when the agent first uses a tool.
    """
    included = get_tools_for_context(role_tag, task_type, all_tool_names)

    if included is None:
        # Default role — include all known summaries
        tool_names = sorted(TOOL_SUMMARIES.keys())
    else:
        tool_names = sorted(included)

    lines = ["Available tools (use any tool to see full documentation):"]
    for name in tool_names:
        summary = TOOL_SUMMARIES.get(name, "")
        if summary:
            lines.append(f"  - {name}: {summary}")
        else:
            lines.append(f"  - {name}")

    return "\n".join(lines)
