import json
from typing import List, Dict, Any

import httpx
from fastapi import FastAPI, HTTPException, Body
from pydantic import BaseModel
from qdrant_client import QdrantClient
from qdrant_client.http.models import PointStruct, VectorParams
from langchain.memory import ConversationBufferMemory
from langchain.embeddings.openai import OpenAIEmbeddings

# ------------------------------------------------------------------
# Конфигурация сервера и клиента
# ------------------------------------------------------------------

QDRANT_HOST = "localhost"
QDRANT_PORT = 6333
COLLECTION_NAME = "agent_memory"

OPENAI_API_KEY = "YOUR_OPENAI_API_KEY"  # замените на свой ключ

# ------------------------------------------------------------------
# Модели данных для API
# ------------------------------------------------------------------

class MemoryItem(BaseModel):
    text: str
    metadata: Dict[str, Any] | None = None


class QueryParams(BaseModel):
    query: str
    top_k: int = 5


# ------------------------------------------------------------------
# Инициализация Qdrant и LangChain
# ------------------------------------------------------------------

client_qdrant = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)

# Создаём коллекцию при первом запуске (если её нет)
if COLLECTION_NAME not in client_qdrant.get_collections().collections:
    client_qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=1536, distance="Cosine"),
    )

embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
memory_buffer = ConversationBufferMemory(memory_key="chat_history", return_messages=True)


# ------------------------------------------------------------------
# FastAPI приложение
# ------------------------------------------------------------------

app = FastAPI(title="MCP Memory Server")


@app.post("/memory/add")
def add_memory(item: MemoryItem):
    """
    Добавляет новый элемент памяти в Qdrant.
    """
    vector = embeddings.embed_query(item.text)
    point_id = client_qdrant.upload_collection(
        collection_name=COLLECTION_NAME,
        points=[PointStruct(id=None, vector=vector, payload=item.metadata or {})],
    ).points[0].id
    return {"status": "ok", "point_id": point_id}


@app.post("/memory/query")
def query_memory(params: QueryParams):
    """
    Выполняет поиск похожих элементов памяти.
    """
    query_vector = embeddings.embed_query(params.query)
    search_result = client_qdrant.search(
        collection_name=COLLECTION_NAME,
        query_vector=query_vector,
        limit=params.top_k,
    )
    results = [
        {"id": hit.id, "score": hit.score, "payload": hit.payload}
        for hit in search_result
    ]
    return {"results": results}


@app.post("/memory/clear")
def clear_memory():
    """
    Очищает всю коллекцию памяти.
    """
    client_qdrant.delete_collection(collection_name=COLLECTION_NAME)
    # Пересоздаём пустую коллекцию
    client_qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=1536, distance="Cosine"),
    )
    return {"status": "cleared"}


# ------------------------------------------------------------------
# Клиентская часть для взаимодействия с сервером
# ------------------------------------------------------------------

class MemoryClient:
    """
    Простой клиент для работы с MCP-сервером памяти.
    """

    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url.rstrip("/")

    def add(self, text: str, metadata: Dict[str, Any] | None = None) -> int:
        payload = {"text": text, "metadata": metadata}
        resp = httpx.post(f"{self.base_url}/memory/add", json=payload)
        if resp.status_code != 200:
            raise RuntimeError(f"Failed to add memory: {resp.text}")
        return resp.json()["point_id"]

    def query(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        payload = {"query": query, "top_k": top_k}
        resp = httpx.post(f"{self.base_url}/memory/query", json=payload)
        if resp.status_code != 200:
            raise RuntimeError(f"Failed to query memory: {resp.text}")
        return resp.json()["results"]

    def clear(self) -> None:
        resp = httpx.post(f"{self.base_url}/memory/clear")
        if resp.status_code != 200:
            raise RuntimeError(f"Failed to clear memory: {resp.text}")


# ------------------------------------------------------------------
# Пример использования (можно запустить как скрипт)
# ------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    from rich.console import Console

    console = Console()
    console.print("[bold green]Запуск MCP Memory Server...[/]")
    uvicorn.run(app, host="0.0.0.0", port=8000)

    # Пример клиента (не будет выполнен при запуске сервера)
    client = MemoryClient()
    point_id = client.add("Привет, как дела?", {"source": "user"})
    console.print(f"Добавлен пункт памяти с id={point_id}")
    results = client.query("привет")
    console.print("[bold]Результаты поиска:[/]")
    for r in results:
        console.print(r)