"""State management routes - /state/* endpoints."""

from __future__ import annotations

from typing import Any

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

from logos_adk.errors import StateError


class StateVersionItem(BaseModel):
    """State version item."""

    version: int
    created_at: str
    agent_name: str
    session_key: str


class StateVersionsResponse(BaseModel):
    """Response for state versions endpoint."""

    items: list[StateVersionItem]


class StateExportResponse(BaseModel):
    """Response for state export endpoint."""

    state: dict[str, Any]
    session_id: str
    version: int | None = None


class StateImportRequest(BaseModel):
    """Request for state import endpoint."""

    state: dict[str, Any]
    session_id: str


class StateImportResponse(BaseModel):
    """Response for state import endpoint."""

    session_id: str


class StateRollbackRequest(BaseModel):
    """Request for state rollback endpoint."""

    session_id: str
    version: int


def create_state_routes(get_state_store, get_agent_name=None) -> APIRouter:
    """Create state management routes."""
    router = APIRouter(prefix="/state", tags=["state"])

    def get_namespace():
        if get_agent_name:
            return get_agent_name()
        return "default"

    @router.get("/versions")
    async def get_versions(session_id: str) -> StateVersionsResponse:
        """Get state versions for a session."""
        store = get_state_store()
        if store is None:
            raise HTTPException(status_code=404, detail="State store not configured")

        namespace = get_namespace()
        items = []

        if hasattr(store, "history"):
            history = store.history(namespace, session_id)
            for entry in history:
                items.append(
                    StateVersionItem(
                        version=int(entry.get("version", 0)),
                        created_at=str(entry.get("saved_at", entry.get("created_at", ""))),
                        agent_name=str(entry.get("agent_name", namespace)),
                        session_key=str(entry.get("session_id", session_id)),
                    )
                )

        return StateVersionsResponse(items=items)

    @router.get("/export")
    async def export_state(session_id: str) -> StateExportResponse:
        """Export state for a session."""
        store = get_state_store()
        if store is None:
            raise HTTPException(status_code=404, detail="State store not configured")

        namespace = get_namespace()
        if hasattr(store, "export_state"):
            exported = store.export_state(namespace, session_id)
            if exported is not None:
                return StateExportResponse(
                    state=dict(exported.get("state", {})),
                    session_id=str(exported.get("session_id", session_id)),
                    version=exported.get("version"),
                )

        state = store.load(namespace, session_id)
        if state is None:
            state_dict: dict[str, Any] = {}
        elif hasattr(state, "model_dump"):
            state_dict = state.model_dump()
        elif hasattr(state, "dict"):
            state_dict = state.dict()
        else:
            state_dict = {}
        return StateExportResponse(state=state_dict, session_id=session_id)

    @router.post("/import")
    async def import_state(request: StateImportRequest) -> StateImportResponse:
        """Import state for a session."""
        store = get_state_store()
        if store is None:
            raise HTTPException(status_code=404, detail="State store not configured")

        namespace = get_namespace()
        if hasattr(store, "import_state"):
            store.import_state(namespace, {"state": request.state}, session_id=request.session_id)
        else:
            from logos_adk.state import AgentState

            state = AgentState(**request.state)
            state.session_id = request.session_id
            store.save(namespace, state)
        return StateImportResponse(session_id=request.session_id)

    @router.post("/rollback")
    async def rollback_state(request: StateRollbackRequest) -> StateExportResponse:
        """Rollback state to a specific version."""
        store = get_state_store()
        if store is None:
            raise HTTPException(status_code=404, detail="State store not configured")

        namespace = get_namespace()
        if hasattr(store, "rollback"):
            try:
                state = store.rollback(namespace, request.session_id, version=request.version)
            except StateError as exc:
                raise HTTPException(status_code=404, detail="Version not found") from exc
        else:
            history = store.history(namespace, request.session_id)
            if request.version > 0 and request.version <= len(history):
                state = history[request.version - 1].get("state", {})
            else:
                raise HTTPException(status_code=404, detail="Version not found")

        state_dict = {}
        if state is not None:
            if hasattr(state, "model_dump"):
                state_dict = state.model_dump()
            elif hasattr(state, "dict"):
                state_dict = state.dict()
            elif isinstance(state, dict):
                state_dict = state

        # Get the new version count after rollback
        history = store.history(namespace, request.session_id)
        new_version = len(history) if history else 1

        return StateExportResponse(
            state=state_dict, session_id=request.session_id, version=new_version
        )

    @router.get("/shards")
    async def get_shards() -> dict[str, Any]:
        """Get state store shards info."""
        store = get_state_store()
        if store is None:
            raise HTTPException(status_code=404, detail="State store not configured")

        if hasattr(store, "describe"):
            info = store.describe()
        else:
            info = {"backend": type(store).__name__}

        return info

    return router
