"""Background Tool Manager — detach running tools from agent loop.

When the user clicks "→ BG" on a running tool call, the agent loop
stops waiting for that tool and continues with a placeholder result.
The tool keeps running in its thread. When it finishes, the result
is injected as a system message in the conversation.

Usage:
    # In _execute_tool_calls: check if a tool was backgrounded
    if BackgroundToolManager.is_backgrounded(tc_id):
        results[tc_id] = (tc, "[Running in background]")
        BackgroundToolManager.register(tc_id, future, conv_id, agent_name, tool_name)

    # Client action: background a tool
    BackgroundToolManager.background(tc_id)

    # Client action: cancel a background tool
    BackgroundToolManager.cancel(tc_id)
"""

import logging
import threading
import time
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

_lock = threading.Lock()
_backgrounded: Dict[str, dict] = {}  # tc_id → task info
_completed: Dict[str, str] = {}  # "conv_id:tc_id" → result (pending agent pickup)
_pending_bg: set = set()  # tc_ids flagged for backgrounding (before registered)


# ── Claude-code tc_id → MCP request_id bridge ───────────────────────────
# The MCP bridge generates its own request_id (a UUID) that is NOT the
# Claude-code tool_use id the UI sees. When the streaming provider emits
# a tool_call SSE event, it enqueues the CC tc_id + args_hash here;
# tool_relay_service consumes it when the matching MCP request arrives
# and stores the association on its _inflight entry. This lets
# background / kill actions (keyed by the UI-visible tc_id) reach the
# actual in-flight MCP request.
#
# CC runs tools in parallel — we match by (conv, agent, tool_name,
# args_hash) with FIFO fallback on hash collisions.
_cc_pending_tcs_lock = threading.Lock()
_cc_pending_tcs: Dict[tuple, list] = {}  # (conv_id, agent_name) → [entry...]
ANY_TOOL = "__pawflow_any_tool__"
ANY_ARGS_HASH = "__pawflow_any_args__"
_PENDING_CC_TTL_SECONDS = 60


def _prune_cc_pending_locked(key: tuple, now: float | None = None) -> None:
    now = time.time() if now is None else now
    queue = _cc_pending_tcs.get(key) or []
    queue = [e for e in queue if now - e.get("ts", 0) <= _PENDING_CC_TTL_SECONDS]
    if queue:
        _cc_pending_tcs[key] = queue
    else:
        _cc_pending_tcs.pop(key, None)


def _args_hash(args) -> str:
    """Stable hash of tool arguments for tc_id ↔ request_id matching."""
    import hashlib
    import json
    try:
        canonical = json.dumps(args, sort_keys=True, default=str, ensure_ascii=False)
    except Exception:
        canonical = repr(args)
    return hashlib.md5(
        canonical.encode("utf-8", errors="replace"),
        usedforsecurity=False,
    ).hexdigest()[:12]


def enqueue_cc_tc(conv_id: str, agent_name: str, tc_id: str,
                   tool_name: str, args_hash: str) -> None:
    """Register a live provider tool_use id awaiting matching in tool_relay."""
    import time as _time
    key = (conv_id, agent_name)

    # Codex/Gemini app-server can issue the MCP execute request before the
    # provider stream surfaces the tool_call item. If the relay already has
    # the request in flight, bind it now instead of queueing a stale entry.
    if tool_name != ANY_TOOL and args_hash != ANY_ARGS_HASH:
        try:
            from services.tool_relay_service import ToolRelayService
            if ToolRelayService.bind_pending_cc_tc(
                    conv_id, agent_name, tc_id, tool_name, args_hash):
                return
        except Exception:
            logger.debug("[bg-tool] late cc_tc bind failed", exc_info=True)

    with _cc_pending_tcs_lock:
        now = _time.time()
        _prune_cc_pending_locked(key, now)
        queue = _cc_pending_tcs.setdefault(key, [])
        queue.append({
            "tc_id": tc_id,
            "tool_name": tool_name,
            "args_hash": args_hash,
            "ts": now,
        })
        _prune_cc_pending_locked(key, now)


def pop_cc_tc(conv_id: str, agent_name: str, tool_name: str,
               args_hash: str) -> str:
    """Find and remove the matching CC tc_id. Returns '' if no match.

    Matches on (tool_name, args_hash) FIFO-within-collision. When two
    parallel calls have identical args, the first pop wins — acceptable
    since BG/kill on one of two truly-identical calls is ambiguous
    anyway.
    """
    key = (conv_id, agent_name)
    with _cc_pending_tcs_lock:
        _prune_cc_pending_locked(key)
        queue = _cc_pending_tcs.get(key)
        if not queue:
            return ""
        for i, entry in enumerate(queue):
            if entry["tool_name"] == tool_name and entry["args_hash"] == args_hash:
                queue.pop(i)
                if not queue:
                    _cc_pending_tcs.pop(key, None)
                return entry["tc_id"]
        # Some ACP providers announce a MCP tool call before exposing the
        # nested PawFlow tool name/arguments. In that case enqueue a scoped
        # FIFO wildcard so the next relay request from the same agent still
        # maps to the UI-visible tool id for BG/kill.
        for i, entry in enumerate(queue):
            if entry["tool_name"] == ANY_TOOL and entry["args_hash"] == ANY_ARGS_HASH:
                queue.pop(i)
                if not queue:
                    _cc_pending_tcs.pop(key, None)
                return entry["tc_id"]
    return ""


def snapshot_cc_pending(conv_id: str, agent_name: str) -> list:
    """Return a copy of the current pending cc_tc queue for (conv, agent).

    Diagnostic only — used by the tool-relay mismatch log to dump what was
    queued at the moment a pop missed, so we can see which (tool_name,
    args_hash) was waiting vs. what the bridge requested.
    """
    key = (conv_id, agent_name)
    with _cc_pending_tcs_lock:
        _prune_cc_pending_locked(key)
        queue = _cc_pending_tcs.get(key) or []
        return [
            {"tool_name": e["tool_name"],
             "args_hash": e["args_hash"],
             "tc_id": e["tc_id"][-12:],
             "age_s": round(time.time() - e["ts"], 2)}
            for e in queue
        ]


def background(tc_id: str) -> bool:
    """Flag a tool_call for backgrounding.

    Tries three paths in order so this works for all providers:
      1. LLM API (Anthropic/OpenAI direct): the agent loop polls
         is_backgrounded(tc_id) in agent_tool_exec.py — just flag
         _pending_bg and it picks up on the next tick.
      2. Already-registered future: we know the tc_id and can set its
         background event directly.
      3. Claude-code via MCP: no direct future, look up the in-flight
         MCP request by CC tc_id and set its background event (handled
         by ToolRelayService.background_by_tc_id).
    """
    with _lock:
        _pending_bg.add(tc_id)
    # Tool-relay path (Claude-code via MCP bridge)
    try:
        from services.tool_relay_service import ToolRelayService
        ToolRelayService.background_by_tc_id(tc_id)
    except Exception as e:
        logger.debug("[bg-tool] tool-relay bg routing skipped: %s", e)
    logger.info("[bg-tool] flagged %s for background", tc_id)
    return True


def is_backgrounded(tc_id: str) -> bool:
    """Check if a tool_call has been flagged for backgrounding."""
    with _lock:
        return tc_id in _pending_bg


def register(tc_id: str, future, conversation_id: str,
             agent_name: str = "", tool_name: str = "",
             is_claude_code: bool = False, user_id: str = ""):
    """Register a backgrounded tool with its running future."""
    with _lock:
        _pending_bg.discard(tc_id)
        _backgrounded[tc_id] = {
            "future": future,
            "conversation_id": conversation_id,
            "agent_name": agent_name,
            "tool_name": tool_name,
            "is_claude_code": is_claude_code,
            "user_id": user_id,
            "started_at": time.time(),
            "status": "running",
            "result": None,
        }
    # Watch the future in a daemon thread
    t = threading.Thread(target=_watch_future, args=(tc_id,), daemon=True)
    t.start()
    logger.info("[bg-tool] registered %s (%s) for conv %s",
                tc_id, tool_name, conversation_id[:8])


def cancel(tc_id: str) -> bool:
    """Cancel a background tool. Injects 'Cancelled' result."""
    with _lock:
        task = _backgrounded.get(tc_id)
        if not task or task["status"] != "running":
            # Maybe it's still pending (not yet registered)
            if tc_id in _pending_bg:
                _pending_bg.discard(tc_id)
                return True
            return False
        task["status"] = "cancelled"
        future = task.get("future")

    # Try to cancel the future (may not work if already running)
    if future and hasattr(future, 'cancel'):
        future.cancel()

    # Inject cancellation result
    _inject_result(tc_id, "[Cancelled by user]", is_cancel=True)
    logger.info("[bg-tool] cancelled %s", tc_id)
    return True


def has_pending(conversation_id: str) -> bool:
    """Check if any bg tasks are still running for this conversation."""
    with _lock:
        return any(
            t["conversation_id"] == conversation_id and t["status"] == "running"
            for t in _backgrounded.values()
        )


def wait_pending(conversation_id: str, timeout: Optional[float] = None,
                 cancel_check=None) -> int:
    """Wait for all pending bg tasks. Returns count completed.

    timeout: optional explicit maximum wait in seconds. When omitted, wait
    until every background task finishes or cancel_check raises.

    cancel_check: callable that raises if the agent was preempted
    (e.g., user sent a new message). Called every second.
    """
    deadline = (time.time() + timeout) if timeout is not None else None
    completed = 0
    while True:
        with _lock:
            pending = [
                tc_id for tc_id, t in _backgrounded.items()
                if t["conversation_id"] == conversation_id and t["status"] == "running"
            ]
        if not pending:
            break
        if deadline is not None:
            remaining = deadline - time.time()
            if remaining <= 0:
                break
            time.sleep(min(1, remaining))
        else:
            time.sleep(1)
        if cancel_check:
            cancel_check()  # raises AgentCancelled if preempted
        for tc_id in pending:
            with _lock:
                t = _backgrounded.get(tc_id)
                if t and t["status"] != "running":
                    completed += 1
    return completed


def list_tasks(conversation_id: str = "") -> List[dict]:
    """List background tasks, optionally filtered by conversation."""
    with _lock:
        tasks = []
        for tc_id, task in _backgrounded.items():
            if conversation_id and task["conversation_id"] != conversation_id:
                continue
            tasks.append({
                "tc_id": tc_id,
                "tool_name": task["tool_name"],
                "status": task["status"],
                "started_at": task["started_at"],
                "duration": time.time() - task["started_at"],
                "agent_name": task["agent_name"],
            })
        return tasks


def pop_completed(conversation_id: str, tc_id: str) -> Optional[str]:
    """Get and remove a completed bg result. Returns None if not ready."""
    key = f"{conversation_id}:{tc_id}"
    with _lock:
        return _completed.pop(key, None)


def cleanup_done(max_age: float = 300):
    """Remove completed/cancelled tasks older than max_age seconds."""
    with _lock:
        now = time.time()
        to_remove = [
            tc_id for tc_id, task in _backgrounded.items()
            if task["status"] in ("done", "cancelled")
            and now - task["started_at"] > max_age
        ]
        for tc_id in to_remove:
            _backgrounded.pop(tc_id, None)


def _periodic_cleanup():
    """Daemon timer: purge old completed tasks every 60s."""
    cleanup_done(max_age=300)
    t = threading.Timer(60, _periodic_cleanup)
    t.daemon = True
    t.start()


# Start periodic cleanup on module import
threading.Timer(60, _periodic_cleanup).start()


def _watch_future(tc_id: str):
    """Wait for a backgrounded tool's future to complete."""
    with _lock:
        task = _backgrounded.get(tc_id)
    if not task:
        return

    future = task.get("future")
    if not future:
        return

    try:
        # Wait for the future (no timeout — it runs until done)
        tc, result_text = future.result()
        with _lock:
            task = _backgrounded.get(tc_id)
            if not task or task["status"] != "running":
                return  # cancelled while waiting
            task["status"] = "done"
            task["result"] = result_text

        _inject_result(tc_id, result_text)
        logger.info("[bg-tool] %s completed: %d chars", tc_id, len(result_text or ""))

    except Exception as e:
        with _lock:
            task = _backgrounded.get(tc_id)
            if task and task["status"] == "running":
                task["status"] = "error"
                task["result"] = str(e)
        _inject_result(tc_id, f"Error: {e}")
        logger.error("[bg-tool] %s failed: %s", tc_id, e)


def _inject_result(tc_id: str, result_text: str, is_cancel: bool = False):
    """Inject the background tool result into the conversation."""
    with _lock:
        task = _backgrounded.get(tc_id)
    if not task:
        return

    conv_id = task["conversation_id"]
    tool_name = task["tool_name"]
    agent_name = task["agent_name"]

    # Publish SSE event
    try:
        from core.conversation_event_bus import ConversationEventBus
        status = "cancelled" if is_cancel else "done"
        ConversationEventBus.instance().publish_event(conv_id, "bg_task_update", {
            "tc_id": tc_id,
            "tool_name": tool_name,
            "status": status,
            "result": result_text[:500],
            "agent_name": agent_name,
        })
    except Exception:
        logging.getLogger(__name__).debug("Ignored exception", exc_info=True)

    # Inject result into the conversation.
    # Strategy depends on provider:
    # - LLM API: replace the placeholder tool_result in the messages array (direct access)
    # - Claude Code: inject as system message (no direct context access)
    _result_content = "[Cancelled by user]" if is_cancel else result_text
    _is_cc = task.get("is_claude_code", False)

    # 1. Write tool_result to transcript (always — visible at reload)
    try:
        from core.conversation_writer import ConversationWriter
        import uuid as _bg_uuid
        from core.llm_client import stamp_message
        tool_msg = stamp_message({
            "role": "tool",
            "content": _result_content[:500],
            "tool_call_id": tc_id,
            "source": {"type": "system", "name": "background"},
        }, conv_id)
        ConversationWriter.for_conversation(conv_id).enqueue_message(tool_msg)
    except Exception as e:
        logger.error("[bg-tool] failed to write tool_result to transcript: %s", e)

    # 2. Store result for agent pickup (agent applies it to in-memory messages
    #    between iterations — no race condition with agent's context writes)
    with _lock:
        _completed[f"{conv_id}:{tc_id}"] = _result_content
    logger.info("[bg-tool] stored result for agent pickup: %s (%d chars)", tc_id, len(_result_content))

    # 3. Claude Code only: also inject system message
    # (CC has its own session — context replacement alone won't reach it,
    # context is only sent on new session which can't happen with bg tools)
    if _is_cc:
        try:
            from core.conversation_writer import ConversationWriter
            if is_cancel:
                content = (
                    f"Your tool call {tc_id} ({tool_name}) was cancelled by "
                    f"the user.\n"
                    f"Partial output (if any):\n{_result_content or '[Cancelled before any output]'}"
                )
            else:
                # Offload large outputs to FileStore so BG injection does
                # not flood CC's context. Keep the first 500 chars inline
                # as a preview, same pattern as ToolRegistry.execute.
                _payload = _result_content or ""
                _offload_threshold = 50_000
                if len(_payload) > _offload_threshold:
                    try:
                        from core.file_store import FileStore
                        _fid = FileStore.instance().store(
                            f"bg_result_{tc_id}.txt",
                            _payload.encode("utf-8", errors="replace"),
                            "text/plain",
                            category="tool_result",
                            user_id=task.get("user_id", "") or "",
                            conversation_id=conv_id,
                            ttl=4 * 3600,
                        )
                        content = (
                            f"Here is the result of your tool call {tc_id} "
                            f"({tool_name}):\n"
                            f"Full output stored at fs://filestore/{_fid}/bg_result_{tc_id}.txt "
                            f"({len(_payload):,} chars).\n"
                            f"Preview:\n{_payload[:500]}\n..."
                        )
                    except Exception as _fe:
                        logger.warning("[bg-tool] FileStore offload failed: %s", _fe)
                        content = (
                            f"Here is the result of your tool call {tc_id} "
                            f"({tool_name}):\n{_payload}"
                        )
                else:
                    content = (
                        f"Here is the result of your tool call {tc_id} "
                        f"({tool_name}):\n{_payload}"
                    )
            from core.llm_client import stamp_message
            msg = stamp_message({
                "role": "user",
                "content": content,
                "source": {"type": "system", "name": "background"},
            }, conv_id)
            # 1. Transcript + shared + other agents' contexts (all routes
            #    covered by append_message based on role+source).
            ConversationWriter.for_conversation(conv_id).enqueue_message(msg)
            # 2. PendingQueue (agent-visible ingress) + wake.
            # Routes through the same single path every other source
            # uses — no more "BG tool injected but agent never woke"
            # special-casing.
            try:
                from core.pending_queue import PendingQueue
                PendingQueue.for_agent(conv_id, agent_name or "").enqueue(
                    msg, source=f"bg_tool:{tool_name}")
            except Exception as _qe:
                logger.warning("[bg-tool] PendingQueue.enqueue failed: %s", _qe)
        except Exception as e:
            logger.error("[bg-tool] failed to inject CC system message: %s", e)

        # 3. Wake the agent (no-op if already running — current turn
        # will drain the queue; otherwise schedule a fresh turn).
        try:
            from tasks.ai.agent_loop import AgentLoopTask
            AgentLoopTask.wake_agent(
                conv_id, agent_name or "",
                reason=f"[bg-tool] CC result injected ({tool_name})",
                user_id=task.get("user_id", "") or "",
            )
        except Exception as e:
            logger.warning("[bg-tool] failed to schedule agent wake: %s", e)

    # Cleanup old tasks periodically
    cleanup_done(max_age=300)
