"""Per-user IcarusEngine lifecycle manager."""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

from icarus.config.permissions import PermissionConfig
from icarus.engine.engine import IcarusEngine
from icarus.runtime.context import UserContext

logger = logging.getLogger(__name__)


@dataclass
class _PoolEntry:
    engine: IcarusEngine
    refcount: int = 0
    last_active: float = field(default_factory=time.monotonic)


class EnginePool:
    """Manages per-user IcarusEngine instances with lazy init and idle teardown."""

    def __init__(
        self,
        server_dir: Path,
        *,
        idle_timeout: float = 300.0,
        sweep_interval: float = 60.0,
    ) -> None:
        self._server_dir = server_dir
        self._idle_timeout = idle_timeout
        self._sweep_interval = sweep_interval
        self._engines: dict[str, _PoolEntry] = {}
        self._lock = asyncio.Lock()
        self._sweep_task: asyncio.Task[None] | None = None
        self._engine_init: Callable[[IcarusEngine, str], None] | None = None

    def set_engine_initializer(self, fn: Callable[[IcarusEngine, str], None]) -> None:
        """Register a callback invoked on each newly-booted engine.

        Used by app.py to install MCP callbacks and thread-state observers
        that require references to the ConnectionManager (passed via closure).
        """
        self._engine_init = fn

    @property
    def server_dir(self) -> Path:
        return self._server_dir

    @property
    def active_users(self) -> list[str]:
        return list(self._engines.keys())

    def _resolve_db_path(
        self, user_id: str, agent_name: str | None = None,
    ) -> Path:
        """Resolve checkpointer DB path. Indirection point for Phase 3."""
        if agent_name:
            return self._server_dir / "agents" / agent_name / "threads.db"
        return self._server_dir / "users" / user_id / "user.db"

    async def _boot_engine(self, engine: IcarusEngine) -> None:
        """Boot an engine instance. Separated for test mocking."""
        await engine.boot()

    async def get_or_create(self, user_id: str) -> IcarusEngine:
        """Get or create an engine for the given user. Bumps refcount."""
        async with self._lock:
            entry = self._engines.get(user_id)
            if entry is not None:
                entry.refcount += 1
                entry.last_active = time.monotonic()
                return entry.engine

            db_path = self._resolve_db_path(user_id)
            user_ctx = UserContext.for_serve_user(user_id, self._server_dir)
            perm_config = PermissionConfig(
                allow_all=False,
                allowed_folders=[user_ctx.home_dir],
                icarus_dir=user_ctx.home_dir,
            )

            # Import here to avoid circular dep
            from icarus.runtime.backend import PermissionedBackend
            from deepagents.backends import LocalShellBackend
            inner = LocalShellBackend()
            backend = PermissionedBackend(inner=inner, config=perm_config)
            engine = IcarusEngine(
                workspace=user_ctx.working_dir,
                db_path=db_path,
                backend=backend,
                user_ctx=user_ctx,
            )

        # Install callbacks BEFORE boot so MCP boot failures are broadcast
        if self._engine_init:
            self._engine_init(engine, user_id)

        # Boot outside the lock to avoid blocking other users
        try:
            await self._boot_engine(engine)
        except Exception:
            await engine.shutdown()
            raise

        async with self._lock:
            # Double-check: another task may have created while we were booting
            if user_id in self._engines:
                await engine.shutdown()
                entry = self._engines[user_id]
                entry.refcount += 1
                return entry.engine
            self._engines[user_id] = _PoolEntry(engine=engine, refcount=1)
            return engine

    def get(self, user_id: str) -> IcarusEngine | None:
        """Get engine without bumping refcount. Returns None if not cached."""
        entry = self._engines.get(user_id)
        return entry.engine if entry else None

    async def release(self, user_id: str) -> None:
        """Decrement refcount. Idle-eligible when refcount hits 0."""
        async with self._lock:
            entry = self._engines.get(user_id)
            if entry is None:
                return
            entry.refcount = max(0, entry.refcount - 1)
            if entry.refcount == 0:
                entry.last_active = time.monotonic()

    def start_sweep(self) -> None:
        """Start the background idle sweep task."""
        if self._sweep_task is None or self._sweep_task.done():
            self._sweep_task = asyncio.create_task(self._idle_sweep())

    async def _idle_sweep(self) -> None:
        """Periodically check for and teardown idle engines."""
        while True:
            await asyncio.sleep(self._sweep_interval)
            now = time.monotonic()
            expired: list[_PoolEntry] = []
            async with self._lock:
                for uid in [
                    uid for uid, e in self._engines.items()
                    if e.refcount == 0 and (now - e.last_active) > self._idle_timeout
                ]:
                    entry = self._engines.pop(uid)
                    expired.append(entry)
            # Shutdown outside the lock (may be slow)
            for entry in expired:
                try:
                    await entry.engine.shutdown()
                except Exception:
                    logger.exception("sweep: failed to shut down engine")

    async def shutdown(self) -> None:
        """Shutdown all engines and stop the sweep task."""
        if self._sweep_task and not self._sweep_task.done():
            self._sweep_task.cancel()
            try:
                await self._sweep_task
            except asyncio.CancelledError:
                pass
        async with self._lock:
            engines = list(self._engines.values())
            self._engines.clear()
        for entry in engines:
            await entry.engine.shutdown()
