import os
from typing import List

from fastmcp import App, Route  # FastMCP – lightweight web framework
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores.qdrant import Qdrant
from langchain.llms.openai import OpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from rich.console import Console

# Конфигурация
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("Не задана переменная окружения OPENAI_API_KEY")

# Инициализация консоли rich
console = Console()

# Векторный хранилище в Qdrant
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
qdrant_store = Qdrant(
    client=None,
    collection_name="rag_memory",
    embeddings=embeddings,
    url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
)

# LLM и цепочка RAG
llm = OpenAI(openai_api_key=OPENAI_API_KEY, temperature=0.7)
prompt_template = PromptTemplate(
    input_variables=["context", "question"],
    template="Ниже приведена информация:\n{context}\n\nВопрос: {question}\nОтвет:",
)
rag_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=qdrant_store.as_retriever(search_kwargs={"k": 5}),
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt_template},
)

app = App()


@app.post("/add")
def add_document(text: str):
    """
    Добавляет новый документ в память агента.
    """
    try:
        qdrant_store.add_texts([text])
        console.log(f"[green]Документ добавлен[/green]")
        return {"status": "ok"}
    except Exception as e:
        console.print_exception()
        return {"status": "error", "detail": str(e)}


@app.post("/ask")
def ask(question: str):
    """
    Делает запрос к агенту с RAG‑памятью.
    """
    try:
        result = rag_chain({"question": question})
        answer = result["answer"]
        sources = [doc.metadata.get("source", "unknown") for doc in result["source_documents"]]
        console.log(f"[blue]Ответ[/blue]: {answer}")
        return {"answer": answer, "sources": sources}
    except Exception as e:
        console.print_exception()
        return {"status": "error", "detail": str(e)}


if __name__ == "__main__":
    # Запуск FastMCP сервера
    app.run(host="0.0.0.0", port=8000)