"""Background task manager — runs detached tool executions for the agent.

A tool call with ``background: true`` (only for tools listed in
``settings.backgroundable_tools``) is submitted here by the ToolExecutor: the
tool's coroutine keeps running in its own asyncio task while the agent turn
continues. The caller immediately gets a ToolResult carrying the task_id.

Lifecycle: running → completed | failed | cancelled. On terminal state the
manager publishes a ``TaskUpdate`` event (out-of-band, via a callback wired to
the orchestrator) and records a pending note that the next turn's
``run_stream`` drains into the session context (see navi/core/task_notes.py).

Tasks are in-memory: a server restart loses them (terminal processes die with
the server anyway). Session Stop does NOT cancel background tasks — only
``tasks cancel`` does; each job carries its own stop_event, independent from
the run's.
"""

import asyncio
import contextvars
import json
import time
import uuid
from collections import deque
from dataclasses import dataclass, field
from datetime import UTC, datetime

import structlog

from navi.tools._internal.base import (
    ToolContext,
    ToolResult,
    current_event_sink,
    current_stop_event,
)
from navi.core.events import SubagentComplete, TaskUpdate

log = structlog.get_logger()


class BoundedEventQueue(asyncio.Queue):
    """Non-blocking ring queue for background-task events.

    ``put`` never waits: when full, the oldest event is dropped. Producers
    (subagent_runner does ``await sink.put(...)``) must never block or raise —
    a detached task with a dead consumer must not leak or stall.
    """

    def __init__(self, maxsize: int = 50) -> None:
        super().__init__(maxsize=maxsize)

    async def put(self, item) -> None:  # noqa: D401 — asyncio.Queue signature
        while self.full():
            try:
                self.get_nowait()
            except asyncio.QueueEmpty:
                break
        self.put_nowait(item)

    def put_nowait(self, item) -> None:
        while self.full():
            try:
                self.get_nowait()
            except asyncio.QueueEmpty:
                break
        super().put_nowait(item)


@dataclass
class TaskJob:
    task_id: str
    session_id: str
    tool: str
    args: dict                          # background flag stripped
    parent_tool_call_id: str = ""
    status: str = "running"             # running | completed | failed | cancelled
    result: ToolResult | None = None
    error: str | None = None
    created_at: float = field(default_factory=time.time)
    finished_at: float | None = None
    ring: BoundedEventQueue | None = None
    done: asyncio.Event = field(default_factory=asyncio.Event)
    stop_event: asyncio.Event = field(default_factory=asyncio.Event)
    task: asyncio.Task | None = None
    subagent_tokens: int | None = None

    def preview(self, limit: int = 800) -> str:
        """Short result summary for notes / task_update events."""
        if self.status == "running":
            return ""
        if self.result is None:
            return self.error or self.status
        text = self.result.to_message_content()
        return text[:limit]

    def to_update(self) -> TaskUpdate:
        return TaskUpdate(
            task_id=self.task_id,
            session_id=self.session_id,
            tool=self.tool,
            status=self.status,
            result_preview=self.preview(),
            parent_tool_call_id=self.parent_tool_call_id,
            started_at=datetime.fromtimestamp(self.created_at, tz=UTC).isoformat(),
            finished_at=(
                datetime.fromtimestamp(self.finished_at, tz=UTC).isoformat()
                if self.finished_at is not None
                else None
            ),
            subagent_tokens=self.subagent_tokens,
        )


class TaskManager:
    """Registry of detached tool executions."""

    FINISHED_CAP_PER_SESSION = 50

    def __init__(self) -> None:
        self._jobs: dict[str, TaskJob] = {}
        self._submit_times: dict[str, deque[float]] = {}  # session_id → timestamps
        self._update_callback = None  # callable(TaskUpdate) | None
        self._sweeper_task: asyncio.Task | None = None

    # -- wiring -------------------------------------------------------------

    def set_update_callback(self, callback) -> None:
        """Inject the orchestrator notifier (called once at container build)."""
        self._update_callback = callback

    def _publish(self, job: TaskJob) -> None:
        if self._update_callback is None:
            return
        try:
            self._update_callback(job.to_update())
        except Exception:
            log.exception("tasks.publish_failed", task_id=job.task_id)

    # -- caps ---------------------------------------------------------------

    def _rate_limited(self, session_id: str) -> bool:
        times = self._submit_times.setdefault(session_id, deque())
        now = time.time()
        while times and now - times[0] > 300:
            times.popleft()
        from navi.config import settings

        return len(times) >= settings.tasks_rate_limit

    def _cap_reason(self, session_id: str, tool: str) -> str | None:
        from navi.config import settings

        running = [j for j in self._jobs.values() if j.status == "running"]
        if len(running) >= settings.tasks_max_global:
            return "global task limit reached"
        per_session = [j for j in running if j.session_id == session_id]
        if len(per_session) >= settings.tasks_max_per_session:
            return "per-session task limit reached"
        if tool == "spawn_agent":
            spawns = [j for j in per_session if j.tool == "spawn_agent"]
            if len(spawns) >= settings.tasks_max_spawn:
                return "background subagent limit reached"
        return None

    # -- API ----------------------------------------------------------------

    def submit(
        self,
        session_id: str,
        tool: str,
        args: dict,
        coro_factory,
        ctx: ToolContext | None,
        parent_tool_call_id: str = "",
        ring_size: int | None = None,
    ) -> TaskJob | str:
        """Detach a tool execution. Returns a TaskJob, or a rejection reason.

        ``coro_factory`` is a one-arg callable receiving the rebuilt background
        ToolContext and returning the coroutine to run (invoked inside the
        detached task, after ContextVar overrides). Never close over the
        caller's own ToolContext.
        """
        if self._rate_limited(session_id):
            return "task rate limit reached (try again later)"
        reason = self._cap_reason(session_id, tool)
        if reason:
            return reason

        from navi.config import settings

        job = TaskJob(
            task_id=f"bt-{uuid.uuid4().hex[:8]}",
            session_id=session_id,
            tool=tool,
            args=args,
            parent_tool_call_id=parent_tool_call_id,
            ring=BoundedEventQueue(maxsize=ring_size or settings.tasks_event_buffer_size),
        )
        job.task = asyncio.create_task(
            self._run_job(job, coro_factory, ctx),
            name=f"bgtask-{job.task_id}",
        )
        self._jobs[job.task_id] = job
        self._submit_times.setdefault(session_id, deque()).append(time.time())
        self._ensure_sweeper()
        self._publish(job)
        log.info(
            "tasks.submitted", task_id=job.task_id, tool=tool,
            session_id=session_id, parent_tool_call_id=parent_tool_call_id,
        )
        return job

    def get(self, task_id: str, session_id: str) -> TaskJob | None:
        job = self._jobs.get(task_id)
        if job is not None and job.session_id == session_id:
            return job
        return None

    def list(self, session_id: str) -> list[TaskJob]:
        return sorted(
            (j for j in self._jobs.values() if j.session_id == session_id),
            key=lambda j: j.created_at,
        )

    def cancel(self, job: TaskJob) -> bool:
        if job.status != "running" or job.task is None or job.task.done():
            return False
        job.stop_event.set()
        job.task.cancel()
        # A task cancelled before its first step never enters _run_job, so its
        # CancelledError handler never fires — finalise by hand via the
        # done-callback (a no-op when _run_job already finalised).
        job.task.add_done_callback(lambda _t: self._finalise_if_unfinished(job))
        return True

    def _finalise_if_unfinished(self, job: TaskJob) -> None:
        if job.status != "running":
            return  # normal path — _run_job already finalised
        job.status = "cancelled"
        job.error = "Task was cancelled."
        job.result = ToolResult(success=False, output="Task was cancelled.",
                                error="cancelled")
        job.finished_at = time.time()
        self._harvest_subagent_tokens(job)
        self._publish(job)
        job.done.set()
        log.info("tasks.finished", task_id=job.task_id, tool=job.tool,
                 status="cancelled", session_id=job.session_id)

        async def _note():
            try:
                from navi.core.task_notes import add_note

                await add_note(job)
            except Exception:
                log.exception("tasks.note_write_failed", task_id=job.task_id)

        asyncio.create_task(_note())

    def running_count(self, session_id: str) -> int:
        return sum(
            1 for j in self._jobs.values()
            if j.session_id == session_id and j.status == "running"
        )

    # -- execution ----------------------------------------------------------

    async def _run_job(self, job: TaskJob, coro_factory, ctx: ToolContext | None) -> None:
        """Run the detached coroutine with per-task sink/stop overrides."""
        from navi.tools._internal.base import (
            current_user_id,
            current_user_role,
            current_user_info,
            current_session_id,
        )

        # Copy the caller's context, then override the run-scoped vars so the
        # detached job (a) streams events into its own bounded ring instead of
        # the dead foreground sink, and (b) ignores the session's Stop signal.
        bg_ctx = ToolContext(
            session_id=getattr(ctx, "session_id", None),
            event_sink=job.ring,
            stop_event=job.stop_event,
            model=getattr(ctx, "model", None),
            user_id=getattr(ctx, "user_id", None),
            user_role=getattr(ctx, "user_role", "user"),
            user_info=getattr(ctx, "user_info", None),
            cwd=getattr(ctx, "cwd", None),
        )
        overrides: list[tuple[contextvars.ContextVar, object]] = [
            (current_event_sink, job.ring),
            (current_stop_event, job.stop_event),
            (current_session_id, job.session_id),
        ]
        if ctx is not None:
            overrides += [
                (current_user_id, getattr(ctx, "user_id", None)),
                (current_user_role, getattr(ctx, "user_role", "user")),
                (current_user_info, getattr(ctx, "user_info", None)),
            ]

        for var, value in overrides:
            var.set(value)
        try:
            job.result = await coro_factory(bg_ctx)
            job.status = "completed"
        except asyncio.CancelledError:
            job.status = "cancelled"
            job.error = "Task was cancelled."
            job.result = ToolResult(success=False, output="Task was cancelled.",
                                    error="cancelled")
        except Exception as e:
            job.status = "failed"
            job.error = f"{type(e).__name__}: {e}"
            job.result = ToolResult(success=False, output=str(e), error=type(e).__name__)
        finally:
            job.finished_at = time.time()
            # ContextVar reset is unnecessary: the task's context dies with it.
            self._harvest_subagent_tokens(job)
            self._publish(job)
            job.done.set()  # never gated on the note write below
            try:
                from navi.core.task_notes import add_note

                await add_note(job)
            except Exception:
                log.exception("tasks.note_write_failed", task_id=job.task_id)
            log.info(
                "tasks.finished", task_id=job.task_id, tool=job.tool,
                status=job.status, session_id=job.session_id,
            )

    def _harvest_subagent_tokens(self, job: TaskJob) -> None:
        """Pull SubagentComplete token counts out of the drained ring."""
        if job.ring is None:
            return
        tokens = None
        while True:
            try:
                item = job.ring.get_nowait()
            except asyncio.QueueEmpty:
                break
            if isinstance(item, SubagentComplete):
                tokens = item.token_count
        job.subagent_tokens = tokens

    # -- maintenance --------------------------------------------------------

    def _ensure_sweeper(self) -> None:
        if self._sweeper_task is None or self._sweeper_task.done():
            self._sweeper_task = asyncio.create_task(self._sweep_loop())

    async def _sweep_loop(self) -> None:
        from navi.config import settings

        while True:
            await asyncio.sleep(60)
            self.reap()

    def reap(self) -> int:
        """Drop finished jobs past TTL / cap. Returns number reaped."""
        from navi.config import settings

        now = time.time()
        reaped = 0
        finished_per_session: dict[str, int] = {}
        for job in list(self._jobs.values()):
            if job.status == "running":
                continue
            finished_per_session[job.session_id] = (
                finished_per_session.get(job.session_id, 0) + 1
            )
            age = now - (job.finished_at or now)
            if age > settings.tasks_ttl_sec:
                del self._jobs[job.task_id]
                reaped += 1
        # Enforce the finished-jobs cap (oldest first) per session.
        for session_id, count in finished_per_session.items():
            if count <= self.FINISHED_CAP_PER_SESSION:
                continue
            finished = sorted(
                (j for j in self._jobs.values()
                 if j.session_id == session_id and j.status != "running"),
                key=lambda j: j.finished_at or 0,
            )
            for job in finished[: count - self.FINISHED_CAP_PER_SESSION]:
                del self._jobs[job.task_id]
                reaped += 1
        return reaped


_manager: TaskManager | None = None


def get_task_manager() -> TaskManager:
    global _manager
    if _manager is None:
        _manager = TaskManager()
    return _manager


def args_summary(args: dict, limit: int = 200) -> str:
    """Compact human-readable argument summary for task ids / notes."""
    try:
        return json.dumps(args, ensure_ascii=False)[:limit]
    except (TypeError, ValueError):
        return str(args)[:limit]