#!/usr/bin/env python3
import os

import requests
from fastmcp import FastMCP
from qdrant_client import QdrantClient

MCP_NAME = os.getenv("MCP_NAME", "local-rag")
QDRANT_URL = os.getenv("QDRANT_URL", "http://127.0.0.1:6333")
COLLECTION = os.getenv("QDRANT_COLLECTION", "docs")
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
GEN_MODEL = os.getenv("GEN_MODEL", "qwen2.5-coder:7b")

mcp = FastMCP(MCP_NAME)
qdrant = QdrantClient(url=QDRANT_URL)


def embed(text: str) -> list[float]:
    r = requests.post(
        "http://127.0.0.1:11434/api/embeddings",
        json={"model": EMBED_MODEL, "prompt": text},
        timeout=120,
    )
    r.raise_for_status()
    return r.json()["embedding"]


def retrieve(query: str, top_k: int = 5):
    query_vec = embed(query)
    return qdrant.search(collection_name=COLLECTION, query_vector=query_vec, limit=top_k)


def sanitize(text: str, max_len: int = 1800) -> str:
    return (text or "").strip()[:max_len]


@mcp.tool
def search_docs(query: str, top_k: int = 5) -> str:
    hits = retrieve(query, top_k=top_k)
    if not hits:
        return "No relevant docs found."
    lines = []
    for i, hit in enumerate(hits, start=1):
        src = hit.payload.get("source", "unknown")
        txt = sanitize(hit.payload.get("text", ""))
        lines.append(f"[{i}] source={src}\n{txt}")
    return "\n\n---\n\n".join(lines)


@mcp.tool
def answer_with_citations(query: str, top_k: int = 5) -> str:
    hits = retrieve(query, top_k=top_k)
    if not hits:
        return "Не нашел релевантных документов."

    context = []
    citations = []
    for i, hit in enumerate(hits, start=1):
        src = hit.payload.get("source", "unknown")
        txt = sanitize(hit.payload.get("text", ""))
        context.append(f"[{i}] SOURCE: {src}\n{txt}")
        citations.append(f"[{i}] {src}")

    prompt = (
        "Ты отвечаешь только на основе контекста.\n"
        "Если данных недостаточно, явно скажи об этом.\n"
        "Добавь ссылки на источники в формате [номер].\n\n"
        f"Вопрос: {query}\n\n"
        "Контекст:\n"
        + "\n\n".join(context)
    )

    r = requests.post(
        "http://127.0.0.1:11434/api/generate",
        json={"model": GEN_MODEL, "prompt": prompt, "stream": False},
        timeout=240,
    )
    r.raise_for_status()
    answer = r.json().get("response", "").strip()

    return f"{answer}\n\nИсточники:\n" + "\n".join(citations)


if __name__ == "__main__":
    mcp.run()