# tools/background_tools.py — Tools for the main agent to manage background tasks.
#
# These expose the SubagentRegistry to the LLM so it can manually dispatch any
# tool to background, check on running tasks, and stop them.
from __future__ import annotations

import json
from langchain_core.tools import tool
from pydantic import BaseModel, Field


class RunInBackgroundInput(BaseModel):
    tool_name: str = Field(description="Name of the tool to run in the background.")
    tool_args: str = Field(
        default="{}",
        description="JSON string of arguments to pass to the tool.",
    )


@tool(args_schema=RunInBackgroundInput)
def run_in_background(tool_name: str, tool_args: str = "{}") -> str:
    """Run any tool as a non-blocking background task. Returns a handle ID immediately.
    The tool runs in a separate thread; use check_background_task to poll status
    and stop_background_task to cancel. Real-time tools (network capture, monitors)
    will stream live data to the dashboard while running in background."""
    from subagent_registry import get_subagent_registry
    from agent_loop import get_openai_tools, get_current_agent_thread_id
    from tool_policy import get_current_policy_context

    registry = get_subagent_registry()
    thread_id = get_current_agent_thread_id() or "unknown"
    tool_name = (tool_name or "").strip()

    _, tools_by_name, full_tools_by_name, _ = get_openai_tools()
    tools_by_name = tools_by_name or full_tools_by_name
    t = tools_by_name.get(tool_name)
    if t is None:
        return f"Unknown tool: {tool_name}. Cannot run in background."

    try:
        args = json.loads(tool_args) if isinstance(tool_args, str) else tool_args
    except json.JSONDecodeError as e:
        return f"Invalid JSON for tool_args: {e}"

    from tools._logging import get_plugin_stream_callback
    handle_id = registry.launch(
        tool_name=tool_name,
        tool_args=args,
        tool_callable=t,
        thread_id=thread_id,
        on_plugin_stream=get_plugin_stream_callback(),
        policy_context=get_current_policy_context(),
    )
    from subagent_registry import _emit_bg_activity

    _emit_bg_activity(
        {
            "source": "background_task",
            "state": "started",
            "summary": f"Queued background task: {tool_name}",
            "thread_id": thread_id,
            "job_id": handle_id,
            "meta": {"tool_name": tool_name},
        }
    )
    return (
        f"Background task started.\n"
        f"Handle ID: {handle_id}\n"
        f"Tool: {tool_name}\n"
        f"Use check_background_task(handle_id='{handle_id}') to check status.\n"
        f"Use stop_background_task(handle_id='{handle_id}') to stop it."
    )


class BackgroundTaskInput(BaseModel):
    handle_id: str = Field(description="Handle ID returned by run_in_background.")


@tool(args_schema=BackgroundTaskInput)
def check_background_task(handle_id: str) -> str:
    """Check the status of a background task. Returns status, elapsed time,
    and result preview if completed."""
    from subagent_registry import get_subagent_registry
    registry = get_subagent_registry()
    info = registry.status(handle_id)
    if "error" in info and "handle_id" not in info:
        return info["error"]
    parts = [
        f"Handle: {info['handle_id']}",
        f"Tool: {info['tool_name']}",
        f"Status: {info['status']}",
        f"Elapsed: {info['elapsed_seconds']}s",
    ]
    if "result_preview" in info:
        parts.append(f"Result preview:\n{info['result_preview']}")
    if "error" in info:
        parts.append(f"Error: {info['error']}")
    return "\n".join(parts)


@tool(args_schema=BackgroundTaskInput)
def stop_background_task(handle_id: str) -> str:
    """Stop a running background task. The tool will receive a stop signal
    and finish gracefully."""
    from subagent_registry import get_subagent_registry
    registry = get_subagent_registry()
    info = registry.stop(handle_id)
    if "error" in info and "handle_id" not in info:
        return info["error"]
    return (
        f"Stop signal sent to {info['handle_id']}.\n"
        f"Tool: {info['tool_name']}\n"
        f"Current status: {info['status']}"
    )


@tool
def list_background_tasks() -> str:
    """List all active and recently completed background tasks."""
    from subagent_registry import get_subagent_registry
    from agent_loop import get_current_agent_thread_id

    registry = get_subagent_registry()
    thread_id = get_current_agent_thread_id()
    tasks = registry.list_active(thread_id=thread_id)

    if not tasks:
        return "No background tasks running or recently completed."

    lines = []
    for t in tasks:
        lines.append(
            f"- {t['handle_id']}: {t['tool_name']} [{t['status']}] ({t['elapsed_seconds']}s)"
        )
    return f"{len(tasks)} background task(s):\n" + "\n".join(lines)
