"""IcarusEngine — transport-agnostic agent lifecycle and streaming.

Never writes to stdout.  Transports (sidecar, HTTP server) wrap the
engine and handle I/O serialization.
"""
from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

from icarus.session.events import EVENT_TYPE_MAP, TextComplete, TextDelta
from icarus.session.thread_state import ThreadState, ThreadStateManager

logger = logging.getLogger(__name__)

MAX_CONCURRENT = 5


@dataclass
class _StreamText:
    """Accumulated stream state for mid-stream history replay."""
    human_content: str
    ai_message_id: str | None = None
    ai_chunks: list[str] = field(default_factory=list)

    @property
    def ai_text(self) -> str:
        return "".join(self.ai_chunks)


class IcarusEngine:
    """Transport-agnostic agent engine.

    Owns lifecycle (boot/shutdown), streaming, HITL bridging, and the
    few handlers that mutate engine state (set_model, set_workspace,
    set_credentials, reload_mcp).  Everything else (thread CRUD, skills,
    plugins, connectors) lives in the transport layer calling core/ directly.
    """

    def __init__(
        self,
        *,
        workspace: Path | None = None,
        db_path: Path | None = None,
        backend: Any | None = None,
        user_ctx: Any | None = None,
    ) -> None:
        self._agent: Any = None
        self._backend: Any = None
        self._model: Any = None
        self._checkpointer: Any = None
        self._checkpointer_cm: Any = None
        self._mcp_manager: Any = None
        self._mcp_tools: list[Any] = []
        self._ctx: Any = user_ctx  # None = boot() creates default
        self._ctx_external: bool = user_ctx is not None  # True = server-provided, don't recreate
        self._workspace: Path = workspace or Path.home()
        self._db_path: Path | None = db_path
        self._ext_backend: Any | None = backend
        self._auto_approve: bool = False

        self._streams: dict[str, asyncio.Task[None]] = {}
        self._hitl_futures: dict[str, asyncio.Future[dict[str, Any]]] = {}
        self._hitl_requests: dict[str, list[dict[str, Any]]] = {}
        self._stream_text: dict[str, _StreamText] = {}
        self._queue: asyncio.Queue[tuple[str, str, Callable]] = asyncio.Queue()
        self._mcp_lock: asyncio.Lock = asyncio.Lock()

        self._state_manager = ThreadStateManager()
        self._known_threads: set[str] = set()
        self._titled_threads: set[str] = set()

        self._pending_model: str | None = None
        self._switch_lock = asyncio.Lock()

    # Properties — transport reads engine state through these.
    agent = property(lambda self: self._agent)
    ctx = property(lambda self: self._ctx)
    checkpointer = property(lambda self: self._checkpointer)
    state_manager = property(lambda self: self._state_manager)
    stream_text = property(lambda self: self._stream_text)
    hitl_requests = property(lambda self: self._hitl_requests)
    hitl_futures = property(lambda self: self._hitl_futures)
    streams = property(lambda self: self._streams)
    known_threads = property(lambda self: self._known_threads)
    titled_threads = property(lambda self: self._titled_threads)
    @property
    def pending_model(self) -> str | None:
        return self._pending_model

    @pending_model.setter
    def pending_model(self, value: str | None) -> None:
        self._pending_model = value
    mcp_manager = property(lambda self: self._mcp_manager)
    mcp_tools = property(lambda self: self._mcp_tools)
    mcp_config_path = property(lambda self: self._mcp_config_path())
    auto_approve = property(lambda self: self._auto_approve)
    workspace = property(lambda self: self._workspace)
    db_path = property(lambda self: self._db_path)

    # ── Lifecycle ────────────────────────────────────────────

    async def boot(self) -> None:
        """Init MCP, checkpointer, and schedule background agent creation."""
        from icarus.config.app_config import settings as _settings
        from icarus.config.model_config import clear_caches
        from icarus.mcp import MCPManager
        from icarus.runtime.context import UserContext
        from icarus.session.threads import get_checkpointer

        _settings.reload_credentials()
        clear_caches()
        if not self._ctx_external:
            self._ctx = UserContext.with_workspace(self._workspace)

        # MCP — non-fatal, errors must never prevent boot
        try:
            self._mcp_manager = MCPManager.from_config_file(
                self._mcp_config_path(),
                on_failure=self._on_mcp_server_failure,
                on_recovery=self._on_mcp_server_recovery,
                on_reconnect=self._on_mcp_server_reconnect,
            )
            if self._mcp_manager:
                self._mcp_tools, failed = await self._mcp_manager.start()
                if failed:
                    await self._on_mcp_boot_failures(failed)
        except (Exception, BaseExceptionGroup):
            logger.warning("MCP startup failed", exc_info=True)
            self._mcp_manager = None
            self._mcp_tools = []

        self._checkpointer_cm = get_checkpointer(db_path=self._db_path)
        self._checkpointer = await self._checkpointer_cm.__aenter__()
        asyncio.create_task(self._boot_agent_background())

    async def _boot_agent_background(self) -> None:
        try:
            await asyncio.to_thread(self._try_create_agent_sync)
        except Exception:
            logger.warning("Background agent creation failed", exc_info=True)

    def _try_create_agent_sync(self) -> None:
        if self._pending_model:  # lazy switch will create the right agent
            return
        self._create_agent()

    async def shutdown(self) -> None:
        """Release resources in reverse order. Never raises."""
        for step in [
            self._cancel_all_streams,
            lambda: self._mcp_manager.stop() if self._mcp_manager else asyncio.sleep(0),
            self._close_model,
            lambda: self._checkpointer_cm.__aexit__(None, None, None) if self._checkpointer_cm else asyncio.sleep(0),
        ]:
            try:
                await step()
            except (Exception, BaseExceptionGroup):
                logger.warning("Shutdown step failed", exc_info=True)

    # ── Agent creation ───────────────────────────────────────

    def _create_agent(self) -> None:
        """Build agent graph. Logs and swallows errors."""
        try:
            from icarus.runtime.builder import create_icarus_agent
            result = create_icarus_agent(
                self._ctx, tools=self._mcp_tools or None,
                checkpointer=self._checkpointer,
            )
            self._apply_agent_result(result)
        except Exception:
            logger.warning("Agent creation failed", exc_info=True)

    def _apply_agent_result(self, result: Any) -> None:
        """Apply builder result, preserving externally-injected backend."""
        self._agent = result.graph
        self._backend = self._ext_backend if self._ext_backend is not None else result.backend
        self._model = result.model

    async def _rebuild_agent(self) -> None:
        """Rebuild agent after workspace/MCP changes."""
        from icarus.runtime.builder import create_icarus_agent
        from icarus.runtime.context import UserContext
        if not self._ctx_external:
            # Desktop/CLI: recreate context with current workspace
            self._ctx = UserContext.with_workspace(self._workspace)
        # Server: keep the externally-provided context (user is sandboxed)
        result = await asyncio.to_thread(
            create_icarus_agent,
            self._ctx, tools=self._mcp_tools or None,
            checkpointer=self._checkpointer,
        )
        await self._apply_model_result(result)

    def require_agent(self) -> None:
        """Raise RuntimeError if agent not ready."""
        if self._agent is not None:
            return
        raise RuntimeError(
            "Agent is not ready -- no model is configured. "
            "Please set your API keys in Settings."
        )

    # ── Stream cancellation ──────────────────────────────────

    async def _cancel_stream(self, thread_id: str) -> None:
        task = self._streams.get(thread_id)
        if task and not task.done():
            task.cancel()
            try:
                await task
            except BaseException:
                pass

    async def _cancel_all_streams(self) -> None:
        await asyncio.gather(
            *(self._cancel_stream(tid) for tid in list(self._streams)),
            return_exceptions=True,
        )
        self._streams.clear()

    # ── Model helpers ────────────────────────────────────────

    async def _close_model(self) -> None:
        if self._model and hasattr(self._model, "aclose"):
            await self._model.aclose()

    async def _apply_model_result(self, result: Any) -> None:
        await self._close_model()
        self._apply_agent_result(result)

    def current_model_spec(self) -> str | None:
        """Pending model spec, or active model from settings."""
        if self._pending_model:
            return self._pending_model
        from icarus.config.app_config import settings
        return f"{settings.model_provider}:{settings.model_name}" if settings.model_name else None

    def _mcp_config_path(self) -> Path:
        if self._ctx is None:
            raise RuntimeError("Engine not booted — call boot() first")
        return self._ctx.config_dir / "mcp.json"

    # ── MCP callbacks (logging only — transports override for I/O) ──

    def _on_mcp_server_failure(self, server_name: str, error_msg: str) -> None:
        logger.warning("MCP server '%s' tool failure: %s", server_name, error_msg)
        if self._mcp_manager:
            self._mcp_manager.mark_degraded(server_name, error_msg)

    def _on_mcp_server_recovery(self, server_name: str) -> None:
        logger.info("MCP server '%s' recovered", server_name)
        if self._mcp_manager:
            self._mcp_manager.mark_recovered(server_name)

    def _on_mcp_server_reconnect(self, server_name: str) -> None:
        logger.info("MCP server '%s' is reachable — scheduling reconnect", server_name)

    async def _on_mcp_boot_failures(self, failed: list[str]) -> None:
        for name in failed:
            logger.warning("MCP server '%s' failed to start", name)

    # ── Streaming ────────────────────────────────────────────

    async def send(
        self, thread_id: str, content: str, *,
        emit: Callable[[dict[str, Any]], None],
    ) -> None:
        """Start a streaming turn. Events flow via the emit callback."""
        # Lazy model switch
        async with self._switch_lock:
            if self._pending_model:
                pending = self._pending_model
                try:
                    result = await self._apply_model_switch(pending)
                    if result.ok:
                        self._pending_model = None
                        await self._apply_model_result(result)
                    elif result.error:
                        self._pending_model = None
                        emit({"type": "model_set", "model": self.current_model_spec()})
                        emit({"type": "error", "message": result.error, "thread_id": thread_id})
                        return
                except Exception as exc:
                    self._pending_model = None
                    logger.exception("Lazy model switch failed for %s", pending)
                    emit({"type": "model_set", "model": self.current_model_spec()})
                    emit({"type": "error", "message": str(exc), "thread_id": thread_id})
                    return

        self.require_agent()
        if not thread_id or not content:
            emit({"type": "error", "message": "send requires thread_id and content"})
            return

        if thread_id not in self._known_threads:
            self._known_threads.add(thread_id)
            emit({"type": "thread_created", "thread_id": thread_id, "title": content[:80]})

        if thread_id in self._streams:
            await self._cancel_stream(thread_id)

        if self._state_manager.get_state(thread_id) == ThreadState.ERROR:
            self._state_manager.set_state(thread_id, ThreadState.IDLE)

        if len(self._streams) < MAX_CONCURRENT:
            self._start_stream(thread_id, content, emit)
        else:
            self._state_manager.set_state(thread_id, ThreadState.QUEUED)
            await self._queue.put((thread_id, content, emit))

    def _start_stream(self, thread_id: str, content: str, emit: Callable) -> None:
        self._state_manager.set_state(thread_id, ThreadState.STREAMING)
        task = asyncio.create_task(self._stream(thread_id, content, emit), name=f"stream-{thread_id}")
        self._streams[thread_id] = task
        task.add_done_callback(lambda _: self._on_stream_done(thread_id, emit))

    async def _stream(self, thread_id: str, content: str, emit: Callable) -> None:
        from icarus.session.stream import stream_turn

        st = _StreamText(human_content=content)
        self._stream_text[thread_id] = st
        try:
            async for event in stream_turn(
                agent=self._agent, user_input=content, thread_id=thread_id,
                assistant_id="default", backend=self._backend,
                auto_approve=self._auto_approve,
                resolve_hitl=self._make_hitl_resolver(thread_id, emit),
            ):
                event.thread_id = thread_id
                if isinstance(event, TextDelta):
                    st.ai_message_id = event.message_id
                    st.ai_chunks.append(event.text)
                emit(self._serialize_event(event))

            self._state_manager.force_idle(thread_id)
            emit({"type": "done", "thread_id": thread_id})

            ai_text = st.ai_text
            if thread_id not in self._titled_threads and ai_text:
                self._titled_threads.add(thread_id)
                asyncio.create_task(
                    self._post_stream(thread_id, st.human_content, ai_text, emit)
                )
        except asyncio.CancelledError:
            self._state_manager.force_idle(thread_id)
            emit({"type": "cancelled", "thread_id": thread_id})
        except Exception as exc:
            logger.exception("Streaming error for thread %s", thread_id)
            self._state_manager.set_state(thread_id, ThreadState.ERROR)
            emit({"type": "error", "message": str(exc), "thread_id": thread_id})
        finally:
            self._stream_text.pop(thread_id, None)

    def _on_stream_done(self, thread_id: str, emit: Callable) -> None:
        self._streams.pop(thread_id, None)
        if not self._queue.empty():
            try:
                next_tid, next_content, next_emit = self._queue.get_nowait()
                self._start_stream(next_tid, next_content, next_emit)
            except asyncio.QueueEmpty:
                pass

    # ── HITL bridge ──────────────────────────────────────────

    def _make_hitl_resolver(self, thread_id: str, emit: Callable) -> Callable:
        async def resolver(action_requests, assistant_id=None):
            actions = [{"name": ar.get("name", ""), "args": ar.get("args", {})}
                       for ar in action_requests]
            self._state_manager.set_state(thread_id, ThreadState.HITL)
            emit({"type": "hitl_request", "action_requests": actions,
                  "assistant_id": assistant_id, "thread_id": thread_id})

            future: asyncio.Future = asyncio.get_running_loop().create_future()
            self._hitl_futures[thread_id] = future
            self._hitl_requests[thread_id] = actions
            try:
                decision = await future
                self._state_manager.set_state(thread_id, ThreadState.STREAMING)
                return decision
            finally:
                self._hitl_futures.pop(thread_id, None)
                self._hitl_requests.pop(thread_id, None)
        return resolver

    def approve(self, thread_id: str, *, auto_approve_all: bool = False) -> None:
        if auto_approve_all:
            self._auto_approve = True
        future = self._hitl_futures.get(thread_id)
        if future and not future.done():
            future.set_result({"type": "auto_approve_all" if auto_approve_all else "approve"})
        else:
            logger.warning("HITL approve for thread %s but no pending request", thread_id)

    def reject(self, thread_id: str) -> None:
        future = self._hitl_futures.get(thread_id)
        if future and not future.done():
            future.set_result({"type": "reject"})
        else:
            logger.warning("HITL reject for thread %s but no pending request", thread_id)

    async def cancel(self, thread_id: str) -> None:
        if thread_id in self._streams:
            await self._cancel_stream(thread_id)
        else:
            self._state_manager.force_idle(thread_id)

    # ── Post-stream (background title generation) ────────────

    async def _post_stream(self, thread_id: str, user_message: str,
                           ai_response: str, emit: Callable) -> None:
        from icarus.session.threads import has_custom_title, rename_thread
        from icarus.session.titles import generate_title
        try:
            if await has_custom_title(thread_id, db_path=self._db_path):
                return
            title = await asyncio.wait_for(
                generate_title(self._model, user_message, ai_response), timeout=30,
            )
            if title:
                await rename_thread(thread_id, title, db_path=self._db_path)
                emit({"type": "thread_renamed", "thread_id": thread_id, "title": title})
        except asyncio.TimeoutError:
            logger.debug("Title generation timed out for %s", thread_id)
        except Exception:
            logger.debug("Title generation failed for %s", thread_id, exc_info=True)

    # ── Lazy model switch ────────────────────────────────────

    async def _apply_model_switch(self, spec: str) -> Any:
        """Run switch_model in a thread (non-blocking)."""
        from icarus.session.model import switch_model
        return await asyncio.to_thread(
            switch_model, spec, ctx=self._ctx, mcp_tools=self._mcp_tools,
            checkpointer=self._checkpointer, auto_approve=self._auto_approve,
        )

    # ── Event serialization ──────────────────────────────────

    @staticmethod
    def _serialize_event(event: Any) -> dict[str, Any]:
        event_type = EVENT_TYPE_MAP.get(type(event))
        if event_type is None:
            return {"type": "unknown", "data": str(event)}
        data = asdict(event)
        data["type"] = event_type
        return data

    # ── State-mutating handlers (need engine internals) ──────

    async def set_model(self, model: str) -> dict[str, Any]:
        """Lazy model switch — actual creation on next send()."""
        if self.current_model_spec() == model:
            return {"type": "model_set", "model": model, "already_active": True}
        self._pending_model = model
        from icarus.config.model_config import save_recent_model
        config_path = Path(self._ctx.config_dir) / "config.toml" if self._ctx else None
        save_recent_model(model, config_path=config_path)
        return {"type": "model_set", "model": model}

    async def set_workspace(self, path: str) -> dict[str, Any]:
        new_path = Path(path)
        if not new_path.is_dir():
            raise ValueError(f"Not a directory: {path}")
        # In server mode (external backend = PermissionedBackend), validate path
        if self._ext_backend is not None:
            from icarus.config.permissions import PermissionConfig
            if hasattr(self._ext_backend, 'config') and isinstance(self._ext_backend.config, PermissionConfig):
                if not self._ext_backend.config.is_path_allowed(str(new_path.resolve())):
                    raise ValueError(f"Access denied: {path}")
        await self._cancel_all_streams()
        self._workspace = new_path.resolve()
        await self._rebuild_agent()
        return {"type": "workspace_set", "path": str(self._workspace)}

    async def set_credentials(
        self, credentials: dict[str, str],
        provider_models: dict[str, list[str]] | None = None,
        provider_endpoints: dict[str, dict[str, str]] | None = None,
    ) -> dict[str, Any]:
        if not credentials and not provider_models and not provider_endpoints:
            raise ValueError("set_credentials requires credentials, provider_models, or provider_endpoints")

        from icarus.config.model_config import clear_caches, save_credentials
        env_path = Path(self._ctx.config_dir) / ".env" if self._ctx else None
        save_credentials(credentials, provider_models, provider_endpoints, env_path=env_path)

        from icarus.config.app_config import settings as _settings
        _settings.reload_credentials()
        clear_caches()

        self._pending_model = None
        await self._close_model()
        self._agent = self._backend = self._model = None
        await asyncio.to_thread(self._try_create_agent_sync)

        # Null out agent if provider lost credentials
        if self._agent is not None:
            from icarus.config.model_config import has_provider_credentials
            if _settings.model_provider and not has_provider_credentials(_settings.model_provider):
                self._agent = self._backend = self._model = None

        return {"type": "credentials_set"}

    async def reload_mcp(self) -> None:
        """Stop MCP, re-read config, restart, rebuild agent."""
        async with self._mcp_lock:
            if self._mcp_manager:
                await self._mcp_manager.stop()
            from icarus.mcp import MCPManager
            self._mcp_manager = MCPManager.from_config_file(
                self._mcp_config_path(),
                on_failure=self._on_mcp_server_failure,
                on_recovery=self._on_mcp_server_recovery,
                on_reconnect=self._on_mcp_server_reconnect,
            )
            if self._mcp_manager:
                self._mcp_tools, failed = await self._mcp_manager.start()
                if failed:
                    await self._on_mcp_boot_failures(failed)
            else:
                self._mcp_tools = []
            await self._rebuild_agent()
