from __future__ import annotations

from typing import TYPE_CHECKING

from fastapi import APIRouter
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from niannian.client import NianNian


class IngestBody(BaseModel):
    text: str
    tags: list[str] = Field(default_factory=list)
    source: str = "http"
    project: str | None = None


class AssistBody(BaseModel):
    task: str
    context: str = ""
    project: str | None = None


class QueryBody(BaseModel):
    question: str
    project: str | None = None


class MemorySpaceBody(BaseModel):
    name: str


def create_memory_router(nn: NianNian) -> APIRouter:
    router = APIRouter(tags=["memory"])

    @router.get("/v1/memory/spaces")
    def list_memory_spaces():
        from niannian.memory.space import read_memory_space_registry

        return read_memory_space_registry(nn._home)

    @router.post("/v1/memory/spaces", status_code=201)
    def create_memory_space(payload: MemorySpaceBody):
        from niannian.memory.space import ensure_memory_space, read_memory_space_registry

        space = ensure_memory_space(nn._home, payload.name)
        reg = read_memory_space_registry(nn._home)
        return {"name": space, "spaces": reg["spaces"], "updated_at": reg.get("updated_at")}

    @router.delete("/v1/memory/spaces/{name}")
    def remove_memory_space(name: str):
        from fastapi import HTTPException

        from niannian.memory.space import MemorySpaceError, delete_memory_space, read_memory_space_registry

        try:
            deleted = delete_memory_space(nn._home, name)
        except MemorySpaceError as exc:
            raise HTTPException(status_code=400, detail={"code": "memory_space_error", "message": str(exc)}) from exc
        reg = read_memory_space_registry(nn._home)
        return {"deleted": deleted, "spaces": reg["spaces"], "updated_at": reg.get("updated_at")}

    @router.post("/v1/ingest")
    def ingest(payload: IngestBody):
        project = payload.project
        if project:
            from niannian.memory.space import ensure_memory_space

            project = ensure_memory_space(nn._home, project)
        atom = nn.ingest(payload.text, tags=payload.tags, source=payload.source, project=project)
        return atom.to_dict()

    @router.post("/v1/query")
    def query(payload: QueryBody):
        return nn.query(payload.question, project=payload.project).to_dict()

    @router.post("/v1/assist")
    def assist(payload: AssistBody):
        return nn.assist(payload.task, context=payload.context, project=payload.project).to_dict()

    @router.post("/v1/dream")
    def dream(project: str | None = None):
        return nn.dream(project=project)

    @router.get("/v1/sync/status")
    def sync_status():
        return nn.sync_status()

    @router.post("/v1/sync/push")
    def sync_push(force: bool = False):
        return nn.sync_push(force=force)

    @router.post("/v1/sync/pull")
    def sync_pull():
        return nn.sync_pull()

    return router
