from fastapi import FastAPI, HTTPException
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_core.vectorstores import InMemoryVectorStore
from pydantic import BaseModel, Field


class KeywordEmbeddings(Embeddings):
    terms = [
        "fastapi",
        "rag",
        "retrieval",
        "context",
        "answer",
        "endpoint",
        "validation",
        "langchain",
    ]

    def _embed(self, text: str) -> list[float]:
        lowered = text.lower()
        return [float(lowered.count(term)) for term in self.terms]

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return [self._embed(text) for text in texts]

    def embed_query(self, text: str) -> list[float]:
        return self._embed(text)


documents = [
    Document(
        page_content=(
            "FastAPI exposes the RAG workflow through a POST endpoint that "
            "accepts a question and returns JSON to the calling application."
        ),
        metadata={"source": "service-design"},
    ),
    Document(
        page_content=(
            "LangChain retrieves context from the vector store before the "
            "answer is assembled, so the response can include source labels."
        ),
        metadata={"source": "retrieval-flow"},
    ),
    Document(
        page_content=(
            "Pydantic request models validate incoming JSON before the RAG "
            "chain runs."
        ),
        metadata={"source": "request-validation"},
    ),
]

vector_store = InMemoryVectorStore.from_documents(documents, KeywordEmbeddings())
retriever = vector_store.as_retriever(search_kwargs={"k": 2})


def format_docs(docs: list[Document]) -> str:
    return "\n".join(
        f"{doc.metadata['source']}: {doc.page_content}" for doc in docs
    )


def answer_from_context(payload: dict) -> dict:
    context = payload["context"]
    sources = [
        line.split(":", 1)[0]
        for line in context.splitlines()
        if ":" in line
    ]
    return {
        "answer": (
            "LangChain retrieved "
            f"{', '.join(sources)} and FastAPI returned the answer as JSON "
            "for the calling application."
        ),
        "sources": sources,
    }


rag_chain = (
    {
        "context": retriever | RunnableLambda(format_docs),
        "question": RunnablePassthrough(),
    }
    | RunnableLambda(answer_from_context)
)


class AskRequest(BaseModel):
    question: str = Field(min_length=3)


class AskResponse(BaseModel):
    answer: str
    sources: list[str]


app = FastAPI(title="LangChain RAG API")


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "documents": len(documents)}


@app.post("/ask", response_model=AskResponse)
def ask(payload: AskRequest) -> dict:
    question = payload.question.strip()
    if not question:
        raise HTTPException(status_code=400, detail="Question is required")
    return rag_chain.invoke(question)