"""Minimal FastAPI-based web UI for small Agentbyte demos."""

from __future__ import annotations

import json
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any
from uuid import uuid4

from agentbyte import InMemorySessionStore
from agentbyte.agents import Agent, AgentResponse
from agentbyte.context import AgentContext
from pydantic import BaseModel

EVENT_STREAM_KEY = "microwebui_event_stream"
RUN_COUNTER_KEY = "microwebui_run_counter"
UI_INDEX = Path(__file__).resolve().parent / "ui" / "index.html"


def _require_fastapi() -> tuple[Any, Any, Any, Any, Any]:
    try:
        from fastapi import FastAPI, HTTPException
        from fastapi.responses import FileResponse
        from fastapi.sse import EventSourceResponse, ServerSentEvent
    except ImportError as exc:
        raise RuntimeError(
            "Agentbyte microwebui requires optional dependencies. "
            "Install with `uv sync --extra webui` or `pip install agentbyte[webui]`."
        ) from exc
    return FastAPI, HTTPException, FileResponse, EventSourceResponse, ServerSentEvent


class ChatRequest(BaseModel):
    """Incoming chat payload."""

    message: str
    session_id: str | None = None
    agent_name: str | None = None
    user_id: str | None = None


class AgentSummary(BaseModel):
    """Compact agent metadata for the discovery dropdown."""

    name: str
    description: str
    default: bool = False


class SessionSummary(BaseModel):
    """Compact session metadata for the session picker."""

    session_id: str
    message_count: int
    created_at: str | None = None
    updated_at: str | None = None
    preview: str | None = None


class SessionEvent(BaseModel):
    """Stored SSE-style event payload for a past run."""

    run_index: int
    run_label: str
    prompt: str
    agent_name: str
    event: dict[str, Any]


def _serialize_stream_item(item: Any) -> str:
    """Serialize Agentbyte stream items to JSON for SSE output."""
    if hasattr(item, "model_dump_json"):
        return item.model_dump_json()

    if hasattr(item, "model_dump"):
        return json.dumps(item.model_dump(mode="json"))

    return json.dumps(
        {
            "event_type": "unknown",
            "content": str(item),
            "python_type": type(item).__name__,
        }
    )


def _parse_stream_item(item: Any) -> dict[str, Any]:
    """Parse a stream item into a JSON-friendly dict."""
    return json.loads(_serialize_stream_item(item))


def _iter_visible_messages(context: AgentContext) -> Iterable[dict[str, Any]]:
    """Yield user/assistant/tool messages with JSON-friendly shape."""
    for message in context.messages:
        if hasattr(message, "model_dump"):
            yield message.model_dump(mode="json")


def _session_preview(context: AgentContext) -> str | None:
    """Return a short preview for a stored session."""
    for message in context.messages:
        role = getattr(message, "role", None)
        content = getattr(message, "content", None)
        if role == "user" and content:
            return str(content)[:72]
    return None


def _stored_event_stream(context: AgentContext) -> list[dict[str, Any]]:
    """Return the stored event stream list from session metadata."""
    events = context.metadata.get(EVENT_STREAM_KEY, [])
    return events if isinstance(events, list) else []


def _next_run_index(context: AgentContext) -> int:
    """Return the next run index for this stored session."""
    current_value = context.metadata.get(RUN_COUNTER_KEY, 0)
    return int(current_value) + 1


class MicroWebUIServer:
    """Serve a small multi-agent web UI from a list of Agent objects."""

    def __init__(self, entities: Sequence[Agent]) -> None:
        if not entities:
            raise ValueError("microwebui requires at least one agent")

        self.entities = list(entities)
        self.entity_lookup = {entity.name: entity for entity in self.entities}
        self.session_store = InMemorySessionStore()

    def default_agent(self) -> Agent:
        """Return the first registered agent as the default selection."""
        return self.entities[0]

    def resolve_agent(self, agent_name: str | None, http_exception: Any) -> Agent:
        """Resolve a requested agent name or fall back to the default agent."""
        if agent_name is None:
            return self.default_agent()

        agent = self.entity_lookup.get(agent_name)
        if agent is None:
            raise http_exception(status_code=404, detail=f"Unknown agent: {agent_name}")
        return agent

    async def stream_agent_events(
        self,
        message: str,
        session_id: str | None,
        agent_name: str | None,
        user_id: str | None,
        http_exception: Any,
        server_sent_event: Any,
    ):
        """Stream agent execution as SSE and persist emitted events in memory."""
        selected_agent = self.resolve_agent(agent_name, http_exception)
        effective_session_id = session_id or str(uuid4())
        context = await self.session_store.get(effective_session_id)
        if context is None:
            context = AgentContext(session_id=effective_session_id, user_id=user_id)
        else:
            context.session_id = effective_session_id
            if user_id is not None:
                context.user_id = user_id

        current_run_index = _next_run_index(context)
        current_run_events: list[SessionEvent] = []

        async for event in selected_agent.run_stream(
            message,
            context=context,
            verbose=True,
            stream_tokens=True,
        ):
            event_payload = _parse_stream_item(event)
            current_run_events.append(
                SessionEvent(
                    run_index=current_run_index,
                    run_label=f"Run {current_run_index}",
                    prompt=message,
                    agent_name=selected_agent.name,
                    event=event_payload,
                )
            )
            if isinstance(event, AgentResponse):
                event.context.metadata[EVENT_STREAM_KEY] = [
                    *_stored_event_stream(event.context),
                    *(stored_event.model_dump(mode="json") for stored_event in current_run_events),
                ]
                event.context.metadata[RUN_COUNTER_KEY] = current_run_index
                await self.session_store.save(effective_session_id, event.context)
            yield server_sent_event(data=event_payload)

    def create_app(self) -> Any:
        """Build the FastAPI application for the microwebui server."""
        FastAPI, HTTPException, FileResponse, EventSourceResponse, ServerSentEvent = (
            _require_fastapi()
        )

        app = FastAPI(title="Agentbyte MicroWebUI")

        @app.post("/chat/stream", response_class=EventSourceResponse)
        async def chat_stream(request: ChatRequest) -> Any:
            async for event in self.stream_agent_events(
                request.message,
                request.session_id,
                request.agent_name,
                request.user_id,
                HTTPException,
                ServerSentEvent,
            ):
                yield event

        @app.get("/api/agents")
        async def list_agents() -> list[AgentSummary]:
            default_name = self.default_agent().name
            return [
                AgentSummary(
                    name=agent.name,
                    description=getattr(agent, "description", "") or "",
                    default=agent.name == default_name,
                )
                for agent in self.entities
            ]

        @app.get("/api/sessions")
        async def list_sessions(user_id: str | None = None) -> list[SessionSummary]:
            session_items = (
                await self.session_store.list_by_user(user_id)
                if user_id
                else await self.session_store.list()
            )
            sessions = [
                SessionSummary(
                    session_id=session_id,
                    message_count=len(context.messages),
                    created_at=context.created_at.isoformat() if context.created_at else None,
                    updated_at=context.updated_at.isoformat() if context.updated_at else None,
                    preview=_session_preview(context),
                )
                for session_id, context in session_items
            ]
            sessions.sort(
                key=lambda session: session.updated_at or session.created_at or "",
                reverse=True,
            )
            return sessions

        @app.get("/api/sessions/{session_id}")
        async def get_session(session_id: str) -> dict[str, Any]:
            context = await self.session_store.get(session_id)
            if context is None:
                raise HTTPException(status_code=404, detail="Session not found")

            return {
                "session_id": session_id,
                "message_count": len(context.messages),
                "context": context.model_dump(mode="json"),
                "visible_messages": list(_iter_visible_messages(context)),
            }

        @app.get("/api/sessions/{session_id}/events")
        async def get_session_events(session_id: str) -> dict[str, Any]:
            context = await self.session_store.get(session_id)
            if context is None:
                raise HTTPException(status_code=404, detail="Session not found")

            events = _stored_event_stream(context)
            return {
                "session_id": session_id,
                "event_count": len(events),
                "events": events,
            }

        @app.get("/")
        async def serve_frontend() -> Any:
            return FileResponse(UI_INDEX)

        return app

    def serve(self, host: str = "127.0.0.1", port: int = 8070) -> None:
        """Run the microwebui application with uvicorn."""
        try:
            import uvicorn
        except ImportError as exc:
            raise RuntimeError(
                "Agentbyte microwebui requires optional dependencies. "
                "Install with `uv sync --extra webui` or `pip install agentbyte[webui]`."
            ) from exc

        print(f"Starting Agentbyte MicroWebUI on http://{host}:{port}/")
        uvicorn.run(self.create_app(), host=host, port=port)


def create_app(entities: Sequence[Agent]) -> Any:
    """Create a FastAPI app for the given list of agents."""
    return MicroWebUIServer(entities).create_app()


def serve(
    entities: Sequence[Agent],
    host: str = "127.0.0.1",
    port: int = 8070,
) -> None:
    """Serve a minimal built-in web UI for the given list of agents."""
    MicroWebUIServer(entities).serve(host=host, port=port)
