"""FastAPI server for the Agentbyte WebUI."""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from pathlib import Path
import webbrowser
from typing import Any

from agentbyte.cancellation_token import CancellationToken

from .execution import ExecutionEngine, _serialize_payload
from .models import (
    CreateSessionRequest,
    HealthResponse,
    RunEntityRequest,
    SessionMessagesResponse,
    StatsResponse,
)
from .registry import EntityRegistry
from .session_store import SessionStore
from .sessions import SessionBindingError, SessionManager


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


def _infer_entity_id(entity: Any, index: int) -> str:
    """Derive a stable in-memory entity id for serve(...)."""
    for candidate in (
        getattr(entity, "name", None),
        getattr(entity, "id", None),
        getattr(getattr(entity, "config", None), "name", None),
    ):
        if isinstance(candidate, str) and candidate:
            return candidate
    return f"entity_{index}"


class AgentbyteWebUIServer:
    """Create and configure the Agentbyte WebUI backend."""

    def __init__(
        self,
        entities_dir: str | None = None,
        enable_cors: bool = True,
        cors_origins: list[str] | None = None,
        session_store: SessionStore | None = None,
    ) -> None:
        self.entities_dir = entities_dir
        self.enable_cors = enable_cors
        self.cors_origins = cors_origins or ["*"]
        self.registry = EntityRegistry(entities_dir)
        self.session_manager = SessionManager(store=session_store)
        self.execution_engine = ExecutionEngine(self.session_manager)

    def create_app(self) -> Any:
        """Build the FastAPI application."""
        FastAPI, HTTPException, CORSMiddleware, EventSourceResponse, StaticFiles = (
            _require_fastapi()
        )

        @asynccontextmanager
        async def lifespan(app: Any):
            yield

        app = FastAPI(
            title="Agentbyte WebUI",
            description="Web interface for Agentbyte agents, workflows, and orchestrators",
            version="0.4.0",
            lifespan=lifespan,
        )

        if self.enable_cors:
            app.add_middleware(
                CORSMiddleware,
                allow_origins=self.cors_origins,
                allow_credentials=True,
                allow_methods=["*"],
                allow_headers=["*"],
            )

        @app.get("/api/health", response_model=HealthResponse)
        async def health_check() -> HealthResponse:
            entities = self.registry.list_entities()
            return HealthResponse(
                status="healthy",
                entities_dir=self.entities_dir,
                entities_count=len(entities),
            )

        @app.get("/api/entities")
        async def list_entities() -> Any:
            return self.registry.list_entities()

        @app.get("/api/entities/{entity_id}")
        async def get_entity(entity_id: str) -> Any:
            entity_info = self.registry.get_entity_info(entity_id)
            if entity_info is None:
                raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
            return entity_info

        @app.post("/api/entities/{entity_id}/run")
        async def run_entity(entity_id: str, request: RunEntityRequest) -> Any:
            entity_obj = self.registry.get_entity_object(entity_id)
            entity_info = self.registry.get_entity_info(entity_id)
            if entity_obj is None or entity_info is None:
                raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")

            messages = self.execution_engine.parse_messages(request.messages)

            try:
                if entity_info.type == "agent":
                    if not messages and not request.approval_responses:
                        raise HTTPException(status_code=400, detail="Messages required for agent execution")
                    response = await self.execution_engine.execute_agent(
                        entity_id,
                        entity_obj,
                        messages=messages,
                        session_id=request.session_id,
                        user_id=request.user_id,
                        approval_responses=request.approval_responses,
                    )
                    return _serialize_payload(response)

                if entity_info.type == "orchestrator":
                    if not messages:
                        raise HTTPException(
                            status_code=400,
                            detail="Messages required for orchestrator execution",
                        )
                    response = await self.execution_engine.execute_orchestrator(
                        entity_id,
                        entity_obj,
                        messages=messages,
                        session_id=request.session_id,
                        user_id=request.user_id,
                    )
                    return _serialize_payload(response)

                if request.input_data is None:
                    raise HTTPException(
                        status_code=400,
                        detail="Input data required for workflow execution",
                    )
                execution = await self.execution_engine.execute_workflow(
                    entity_id,
                    entity_obj,
                    input_data=request.input_data,
                    session_id=request.session_id,
                    user_id=request.user_id,
                )
                return _serialize_payload(execution)
            except SessionBindingError as exc:
                raise HTTPException(status_code=409, detail=str(exc)) from exc

        @app.post("/api/entities/{entity_id}/run/stream")
        async def run_entity_stream(entity_id: str, request: RunEntityRequest) -> Any:
            entity_obj = self.registry.get_entity_object(entity_id)
            entity_info = self.registry.get_entity_info(entity_id)
            if entity_obj is None or entity_info is None:
                raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")

            cancellation_token = CancellationToken()
            messages = self.execution_engine.parse_messages(request.messages)
            if entity_info.type == "agent":
                if not messages and not request.approval_responses:
                    raise HTTPException(status_code=400, detail="Messages required for agent execution")
            elif entity_info.type == "orchestrator":
                if not messages:
                    raise HTTPException(status_code=400, detail="Messages required for orchestrator execution")
            elif request.input_data is None:
                raise HTTPException(status_code=400, detail="Input data required for workflow execution")
            try:
                await self.session_manager.assert_compatible(
                    request.session_id,
                    entity_id,
                    user_id=request.user_id,
                )
            except SessionBindingError as exc:
                raise HTTPException(status_code=409, detail=str(exc)) from exc

            async def cancellable_generator() -> Any:
                try:
                    if entity_info.type == "agent":
                        generator = self.execution_engine.execute_agent_stream(
                            entity_id,
                            entity_obj,
                            messages=messages,
                            session_id=request.session_id,
                            user_id=request.user_id,
                            stream_tokens=request.stream_tokens,
                            approval_responses=request.approval_responses,
                            cancellation_token=cancellation_token,
                        )
                    elif entity_info.type == "orchestrator":
                        generator = self.execution_engine.execute_orchestrator_stream(
                            entity_id,
                            entity_obj,
                            messages=messages,
                            session_id=request.session_id,
                            user_id=request.user_id,
                            cancellation_token=cancellation_token,
                        )
                    else:
                        generator = self.execution_engine.execute_workflow_stream(
                            entity_id,
                            entity_obj,
                            input_data=request.input_data,
                            session_id=request.session_id,
                            user_id=request.user_id,
                            cancellation_token=cancellation_token,
                        )

                    async for event in generator:
                        yield event
                except (GeneratorExit, asyncio.CancelledError):
                    cancellation_token.cancel()
                    raise

            return EventSourceResponse(cancellable_generator())

        @app.get("/api/sessions")
        async def list_sessions(
            entity_id: str | None = None,
            user_id: str | None = None,
        ) -> Any:
            return await self.session_manager.list(entity_id=entity_id, user_id=user_id)

        @app.delete("/api/entities/{entity_id}")
        async def delete_entity(entity_id: str) -> Any:
            entity_info = self.registry.get_entity_info(entity_id)
            if entity_info is None:
                raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
            removed = self.registry.unregister_entity(entity_id)
            if not removed:
                raise HTTPException(
                    status_code=404,
                    detail=f"Entity {entity_id} not found or cannot be removed",
                )
            return {
                "status": "success",
                "entity_id": entity_id,
                "message": "Entity removed successfully",
            }

        @app.post("/api/sessions")
        async def create_session(request: CreateSessionRequest) -> Any:
            session_id = self.session_manager.create_session_id()
            context = await self.session_manager.get_or_create(
                session_id,
                request.entity_id,
                request.entity_type,
                user_id=request.user_id,
            )
            return {
                "id": session_id,
                "entity_id": request.entity_id,
                "entity_type": request.entity_type,
                "created_at": context.created_at.isoformat(),
                "last_activity": (context.updated_at or context.created_at).isoformat(),
                "message_count": len(context.messages),
                "user_id": context.user_id,
            }

        @app.get("/api/sessions/{session_id}")
        async def get_session(session_id: str) -> Any:
            context = await self.session_manager.get(session_id)
            if context is None:
                raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
            return context.model_dump(mode="json")

        @app.get("/api/sessions/{session_id}/messages", response_model=SessionMessagesResponse)
        async def get_session_messages(session_id: str) -> SessionMessagesResponse:
            context = await self.session_manager.get(session_id)
            if context is None:
                raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
            metadata = context.metadata if isinstance(context.metadata, dict) else {}
            final_result = metadata.get("final_result")
            stop_reason = metadata.get("stop_reason")
            return SessionMessagesResponse(
                session_id=session_id,
                messages=_serialize_payload(context.messages),
                final_result=final_result if isinstance(final_result, str) else None,
                stop_reason=stop_reason if isinstance(stop_reason, str) else None,
            )

        @app.delete("/api/sessions/{session_id}")
        async def delete_session(session_id: str) -> Any:
            deleted = await self.session_manager.delete(session_id)
            if not deleted:
                raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
            return {"status": "deleted", "session_id": session_id}

        @app.get("/api/stats", response_model=StatsResponse)
        async def get_stats() -> StatsResponse:
            entities = self.registry.list_entities()
            sessions = await self.session_manager.list()
            return StatsResponse(
                entities={
                    "total": len(entities),
                    "by_type": {
                        "agents": len([entity for entity in entities if entity.type == "agent"]),
                        "orchestrators": len([entity for entity in entities if entity.type == "orchestrator"]),
                        "workflows": len([entity for entity in entities if entity.type == "workflow"]),
                    },
                },
                sessions={
                    "total_sessions": len(sessions),
                    "total_messages": sum(session.message_count for session in sessions),
                },
            )

        @app.post("/api/cache/clear")
        async def clear_cache() -> Any:
            self.registry.clear_cache()
            return {"status": "cache_cleared"}

        self._mount_frontend(app, StaticFiles)
        return app

    def _mount_frontend(self, app: Any, StaticFiles: Any) -> None:
        module_dir = Path(__file__).parent
        frontend_dir = module_dir / "ui"
        if frontend_dir.exists() and frontend_dir.is_dir():
            app.mount(
                "/",
                StaticFiles(directory=str(frontend_dir), html=True),
                name="frontend",
            )
            return

        @app.get("/")
        async def root() -> dict[str, str]:
            return {
                "message": "Agentbyte WebUI API",
                "docs": "/docs",
                "health": "/api/health",
                "note": "Frontend assets were not packaged. Rebuild the webui assets before release.",
            }


class WebUIServer:
    """Programmatic server wrapper for registration and startup."""

    def __init__(
        self,
        entities_dir: str | None = None,
        port: int = 8080,
        host: str = "127.0.0.1",
        enable_cors: bool = True,
        cors_origins: list[str] | None = None,
        session_store: SessionStore | None = None,
    ) -> None:
        self._server = AgentbyteWebUIServer(
            entities_dir=entities_dir,
            enable_cors=enable_cors,
            cors_origins=cors_origins,
            session_store=session_store,
        )
        self.port = port
        self.host = host
        self._app: Any | None = None

    def register_entity(self, entity_id: str, entity_obj: Any) -> Any:
        return self._server.registry.register_entity(entity_id, entity_obj)

    def get_app(self) -> Any:
        if self._app is None:
            self._app = self._server.create_app()
        return self._app

    def start(
        self,
        auto_open: bool = False,
        reload: bool = False,
        log_level: str = "info",
    ) -> None:
        webui(
            entities_dir=None,
            port=self.port,
            host=self.host,
            auto_open=auto_open,
            reload=reload,
            log_level=log_level,
            app=self.get_app(),
        )


def create_app(
    entities_dir: str | None = None,
    session_store: SessionStore | None = None,
    **kwargs: Any,
) -> Any:
    """Create a FastAPI app for the WebUI."""
    server = AgentbyteWebUIServer(
        entities_dir=entities_dir,
        session_store=session_store,
        **kwargs,
    )
    return server.create_app()


def webui(
    entities_dir: str | None = None,
    port: int = 8080,
    host: str = "127.0.0.1",
    auto_open: bool = True,
    reload: bool = False,
    log_level: str = "info",
    session_store: SessionStore | None = None,
    app: Any | None = None,
) -> None:
    """Launch the Agentbyte WebUI via Uvicorn."""
    try:
        import uvicorn
    except ImportError as exc:
        raise RuntimeError(
            "Agentbyte WebUI requires optional dependencies. "
            "Install with `uv sync --extra webui` or `pip install agentbyte[webui]`."
        ) from exc

    if app is None:
        app = create_app(entities_dir=entities_dir, session_store=session_store)

    if auto_open:
        import threading
        import time

        def _open() -> None:
            time.sleep(1.2)
            webbrowser.open(f"http://{host}:{port}")

        threading.Thread(target=_open, daemon=True).start()

    uvicorn.run(
        app,
        host=host,
        port=port,
        reload=reload,
        log_level=log_level,
        access_log=True,
    )


def serve(
    entities: list[Any] | None = None,
    entities_dir: str | None = None,
    port: int = 8080,
    host: str = "127.0.0.1",
    auto_open: bool = True,
    session_store: SessionStore | None = None,
) -> None:
    """Serve in-memory and/or discovered entities in the WebUI."""
    server = WebUIServer(
        entities_dir=entities_dir,
        port=port,
        host=host,
        session_store=session_store,
    )
    for index, entity in enumerate(entities or []):
        entity_id = _infer_entity_id(entity, index)
        server.register_entity(entity_id, entity)
    server.start(auto_open=auto_open)


launch = webui


__all__ = [
    "AgentbyteWebUIServer",
    "WebUIServer",
    "create_app",
    "launch",
    "serve",
    "webui",
]
