# subagent_registry.py — Managed background subagents with start/status/stop/list.
#
# Any tool marked realtime=True (or manually dispatched via run_in_background)
# is executed inside a background subagent so the main agent stays responsive.
# Each subagent gets its own stop flag and streams events through the shared
# plugin stream bus.
from __future__ import annotations

import atexit
import contextvars
import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable
from uuid import uuid4

log = logging.getLogger(__name__)

_DEFAULT_MAX_BACKGROUND = 5


def _emit_bg_activity(event: dict, *, tags: list[str] | None = None) -> None:
    """Push background-task lifecycle into the activity bus; never fail the worker."""
    try:
        from activity_bus import emit_activity_enriched

        emit_activity_enriched(event, tags=tags or ["background"])
    except Exception:
        log.debug("background_task activity emit failed", exc_info=True)

try:
    from tools._logging import set_plugin_stream_callback, set_tool_stop_check_callback
except Exception:  # pragma: no cover - defensive fallback for partial startup states
    def set_plugin_stream_callback(_cb):
        return None

    def set_tool_stop_check_callback(_cb):
        return None


@dataclass
class SubagentHandle:
    handle_id: str
    thread_id: str
    tool_name: str
    tool_args: dict
    started_at: float
    status: str = "running"              # running | completed | error | stopped
    result: str | None = None
    error: str | None = None
    _stop_flag: threading.Event = field(default_factory=threading.Event)
    _thread: threading.Thread | None = None


class SubagentRegistry:
    """Thread-safe registry of background tool executions."""

    def __init__(self, max_concurrent: int = _DEFAULT_MAX_BACKGROUND) -> None:
        self._handles: dict[str, SubagentHandle] = {}
        self._lock = threading.Lock()
        self._semaphore = threading.Semaphore(max_concurrent)

    # ── public API ────────────────────────────────────────────────────────

    def launch(
        self,
        tool_name: str,
        tool_args: dict,
        tool_callable: Any,
        thread_id: str,
        *,
        on_plugin_stream: Callable[[dict], None] | None = None,
        policy_context: Any | None = None,
    ) -> str:
        """Start *tool_callable* in a daemon thread. Returns handle_id immediately."""
        self.cleanup()

        if not self._semaphore.acquire(timeout=10):
            raise RuntimeError(f"Too many background tasks running (limit reached). Cannot launch {tool_name}.")

        handle_id = f"bg_{uuid4().hex[:10]}"
        normalized_args = self._normalize_tool_args_for_launch(tool_name, tool_args)
        handle = SubagentHandle(
            handle_id=handle_id,
            thread_id=thread_id,
            tool_name=tool_name,
            tool_args=normalized_args,
            started_at=time.time(),
        )
        launch_ctx = contextvars.copy_context()
        _emit_bg_activity(
            {
                "source": "background_task",
                "state": "started",
                "summary": f"Background task started: {tool_name}",
                "thread_id": thread_id,
                "job_id": handle_id,
                "meta": {"tool_name": tool_name},
            }
        )

        def _run() -> None:
            stream_enabled = on_plugin_stream is not None
            if stream_enabled:
                from tools.job_stream_schema import (
                    open_job_stream,
                    update_job_stream,
                    close_job_stream,
                    error_job_stream,
                    make_capabilities,
                )
                from tools.plugin_stream_controls import clear_stream_state

            sid = handle_id
            try:
                _args_preview = str(handle.tool_args)
                if len(_args_preview) > 300:
                    _args_preview = _args_preview[:300] + "…"
                if on_plugin_stream:
                    set_plugin_stream_callback(on_plugin_stream)
                set_tool_stop_check_callback(lambda: handle._stop_flag.is_set())

                if stream_enabled:
                    sid = open_job_stream(
                        plugin_id="background_tool",
                        plugin_name="Background Tool",
                        title=f"Background: {tool_name}",
                        job_type="background_task",
                        thread_id=thread_id,
                        stream_id=handle_id,
                        controls=[
                            {
                                "id": "stop_task",
                                "type": "button",
                                "label": "Stop Task",
                                "action_id": "stop_task",
                                "style": "destructive",
                            }
                        ],
                        capabilities=make_capabilities(cancel=True),
                        metadata={"tool_name": tool_name, "args_preview": _args_preview},
                    )
                    update_job_stream(
                        plugin_id="background_tool",
                        plugin_name="Background Tool",
                        stream_id=sid,
                        thread_id=thread_id,
                        status="running",
                        phase="start",
                        progress_pct=1.0,
                        message=f"Started background task '{tool_name}'.",
                        extra={"tool_name": tool_name},
                    )

                result = None
                # Policy gate is opt-in for this registry path. In unit tests and
                # internal callers that do not provide policy_context, background
                # execution should mirror direct tool invocation behavior.
                _bg_block = None
                if policy_context is not None:
                    from tool_policy import POLICY_BLOCKED_PREFIX, policy_gate_tool_call

                    try:
                        _bg_block = policy_gate_tool_call(tool_name, handle.tool_args, ctx=policy_context)
                    except Exception:
                        _bg_block = (
                            f"{POLICY_BLOCKED_PREFIX} Policy check failed for background task '{tool_name}'. "
                            "The tool was not executed."
                        )
                if _bg_block is not None:
                    result_str = _bg_block
                    if stream_enabled:
                        close_job_stream(
                            plugin_id="background_tool",
                            plugin_name="Background Tool",
                            stream_id=sid,
                            thread_id=thread_id,
                            status="completed",
                            message=f"Background task '{tool_name}' blocked by policy.",
                            extra={"phase": "bg_done", "tool_name": tool_name, "result_preview": result_str[:500]},
                        )
                    with self._lock:
                        handle.status = "completed"
                        handle.result = result_str
                    return

                result = tool_callable.invoke(handle.tool_args)
                result_str = result if isinstance(result, str) else str(result)
                was_stopped = bool(handle._stop_flag.is_set())

                # Publish terminal state before any stream-close work so status()
                # reflects completion promptly (avoids timing races in callers/tests).
                with self._lock:
                    handle.status = "stopped" if was_stopped else "completed"
                    handle.result = result_str

                if stream_enabled:
                    close_job_stream(
                        plugin_id="background_tool",
                        plugin_name="Background Tool",
                        stream_id=sid,
                        thread_id=thread_id,
                        status="stopped" if was_stopped else "completed",
                        message=(
                            f"Background task '{tool_name}' stopped."
                            if was_stopped
                            else f"Background task '{tool_name}' completed."
                        ),
                        extra={"phase": "bg_done", "tool_name": tool_name, "result_preview": result_str[:500]},
                    )

                _emit_bg_activity(
                    {
                        "source": "background_task",
                        "state": "stopped" if was_stopped else "completed",
                        "summary": (
                            f"Background task stopped: {tool_name}"
                            if was_stopped
                            else f"Background task completed: {tool_name}"
                        ),
                        "thread_id": thread_id,
                        "job_id": handle_id,
                    }
                )

            except Exception as exc:
                if stream_enabled:
                    error_job_stream(
                        plugin_id="background_tool",
                        plugin_name="Background Tool",
                        stream_id=sid,
                        thread_id=thread_id,
                        phase="bg_error",
                        message=str(exc),
                        extra={"tool_name": tool_name},
                    )
                with self._lock:
                    handle.status = "error"
                    handle.error = str(exc)
                _emit_bg_activity(
                    {
                        "source": "background_task",
                        "state": "error",
                        "summary": f"Background task failed: {tool_name}",
                        "thread_id": thread_id,
                        "job_id": handle_id,
                        "meta": {"error": str(exc)[:300]},
                    }
                )

            finally:
                if stream_enabled:
                    clear_stream_state(sid)
                set_plugin_stream_callback(None)
                set_tool_stop_check_callback(None)
                self._semaphore.release()

        t = threading.Thread(
            target=lambda: launch_ctx.run(_run),
            daemon=True,
            name=f"bg-{tool_name}-{handle_id}",
        )
        handle._thread = t

        with self._lock:
            self._handles[handle_id] = handle

        t.start()
        return handle_id

    def status(self, handle_id: str) -> dict:
        """Return current state of a background task."""
        h = self._get_handle(handle_id)
        if h is None:
            return {"error": f"Unknown handle: {handle_id}"}
        elapsed = time.time() - h.started_at
        info: dict[str, Any] = {
            "handle_id": h.handle_id,
            "tool_name": h.tool_name,
            "status": h.status,
            "elapsed_seconds": round(elapsed, 1),
        }
        if h.status == "completed" and h.result is not None:
            info["result_preview"] = h.result[:1000]
        if h.status == "error" and h.error is not None:
            info["error"] = h.error
        return info

    def stop(self, handle_id: str) -> dict:
        """Signal a background task to stop. Returns status dict."""
        h = self._get_handle(handle_id)
        if h is None:
            return {"error": f"Unknown handle: {handle_id}"}
        if h.status != "running":
            return self.status(handle_id)
        h._stop_flag.set()
        _emit_bg_activity(
            {
                "source": "background_task",
                "state": "info",
                "summary": f"Stop requested: {h.tool_name}",
                "thread_id": h.thread_id,
                "job_id": h.handle_id,
            }
        )
        if h._thread is not None:
            h._thread.join(timeout=10)
        return self.status(handle_id)

    def get_result(self, handle_id: str) -> str | None:
        """Return the full result string, or None if not yet available."""
        h = self._get_handle(handle_id)
        if h is None:
            return None
        return h.result

    def list_active(self, thread_id: str | None = None) -> list[dict]:
        """List running (and recently completed) background tasks."""
        with self._lock:
            handles = list(self._handles.values())
        out = []
        for h in handles:
            if thread_id and h.thread_id != thread_id:
                continue
            out.append({
                "handle_id": h.handle_id,
                "tool_name": h.tool_name,
                "status": h.status,
                "elapsed_seconds": round(time.time() - h.started_at, 1),
            })
        return out

    def cleanup(self, max_age_seconds: float = 600) -> int:
        """Remove completed/error handles older than max_age_seconds. Returns count removed."""
        now = time.time()
        to_remove = []
        with self._lock:
            for hid, h in self._handles.items():
                if h.status in ("completed", "error", "stopped") and (now - h.started_at) > max_age_seconds:
                    to_remove.append(hid)
            for hid in to_remove:
                del self._handles[hid]
        return len(to_remove)

    def shutdown_all(self) -> None:
        """Signal all running tasks to stop and join threads with a grace period."""
        with self._lock:
            running = [h for h in self._handles.values() if h.status == "running"]
        for h in running:
            h._stop_flag.set()
        for h in running:
            if h._thread is not None:
                h._thread.join(timeout=5)

    def _get_handle(self, handle_id: str) -> SubagentHandle | None:
        from utils.llm_token_normalize import normalize_bg_handle_token

        raw = str(handle_id or "").strip()
        canonical = normalize_bg_handle_token(raw)
        with self._lock:
            if raw in self._handles:
                return self._handles.get(raw)
            if canonical in self._handles:
                return self._handles.get(canonical)
            for hid, handle in self._handles.items():
                if normalize_bg_handle_token(hid) == canonical:
                    return handle
        return None

    def _normalize_tool_args_for_launch(self, tool_name: str, tool_args: dict) -> dict:
        from tool_arg_normalizer import normalize_tool_args

        args = normalize_tool_args(tool_name, dict(tool_args or {}))

        # If the background tool is run_in_background, normalize nested inner args too.
        if tool_name == "run_in_background":
            inner_name = str(args.get("tool_name") or "").strip()
            inner_raw = args.get("tool_args")
            if inner_name and isinstance(inner_raw, str):
                try:
                    import json
                    parsed = json.loads(inner_raw)
                    if isinstance(parsed, dict):
                        parsed = normalize_tool_args(inner_name, parsed)
                        args["tool_args"] = json.dumps(parsed)
                except Exception:
                    pass
        return args


# Module-level singleton used by tools and the agent loop.
_registry: SubagentRegistry | None = None
_registry_lock = threading.Lock()


def get_subagent_registry() -> SubagentRegistry:
    """Return (and lazily create) the global SubagentRegistry singleton."""
    global _registry
    if _registry is None:
        with _registry_lock:
            if _registry is None:
                try:
                    from core import get_config
                    cfg = (get_config().get("subagent") or {})
                except Exception:
                    cfg = {}
                max_bg = int(cfg.get("max_background", _DEFAULT_MAX_BACKGROUND))
                _registry = SubagentRegistry(max_concurrent=max_bg)
                atexit.register(_registry.shutdown_all)
    return _registry
