from __future__ import annotations

from typing import TYPE_CHECKING, Any

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from niannian.memory.fragment import FragmentValidationError

if TYPE_CHECKING:
    from niannian.client import NianNian


class FragmentBody(BaseModel):
    protocol: str = "niannian.fragment/2"
    project: str = "global"
    content: dict[str, Any]
    taxonomy: dict[str, Any] = Field(default_factory=dict)
    scope: dict[str, Any] = Field(default_factory=dict)
    source: dict[str, Any] = Field(default_factory=dict)
    relations: dict[str, Any] = Field(default_factory=dict)
    retrieval: dict[str, Any] = Field(default_factory=dict)
    dedup_key: str | None = None
    occurred_at: str | None = None

    model_config = {"extra": "allow"}


class FragmentBatchBody(BaseModel):
    fragments: list[FragmentBody]


class RecallBody(BaseModel):
    project: str = "global"
    task: str = ""
    context: str = ""
    scope: dict[str, Any] = Field(default_factory=dict)
    policy: str | None = None
    filters: dict[str, Any] = Field(default_factory=dict)
    return_: dict[str, Any] | None = Field(default=None, alias="return")

    model_config = {"populate_by_name": True}


class ReflectBody(BaseModel):
    project: str | None = None
    scope: dict[str, Any] = Field(default_factory=dict)


def create_mrrr_router(nn: NianNian) -> APIRouter:
    router = APIRouter(tags=["mrrr"])

    @router.post("/v1/fragments")
    def post_fragment(payload: FragmentBody):
        try:
            result = nn.record(fragment=payload.model_dump(exclude_none=True))
        except FragmentValidationError as exc:
            raise HTTPException(status_code=400, detail={"code": "invalid_fragment", "message": str(exc)}) from exc
        return result.to_dict()

    @router.post("/v1/fragments/batch")
    def post_fragments_batch(payload: FragmentBatchBody):
        try:
            results = nn._record.record_batch([f.model_dump(exclude_none=True) for f in payload.fragments])
        except FragmentValidationError as exc:
            raise HTTPException(status_code=400, detail={"code": "invalid_fragment", "message": str(exc)}) from exc
        return {"results": [r.to_dict() for r in results]}

    @router.post("/v1/recall")
    def post_recall(payload: RecallBody):
        return_fmt = "pack"
        if payload.return_ and payload.return_.get("format"):
            return_fmt = str(payload.return_["format"])
        result = nn.recall(
            payload.task,
            context=payload.context,
            scope=payload.scope,
            policy=payload.policy,
            project=payload.project,
            return_format=return_fmt,
            filters=payload.filters or None,
        )
        if hasattr(result, "to_dict"):
            return result.to_dict()
        if isinstance(result, str):
            return {"context_block": result}
        return {"atoms": result}

    @router.post("/v1/reflect")
    def post_reflect(payload: ReflectBody):
        return nn.reflect(scope=payload.scope, project=payload.project).to_dict()

    return router
