import json
import uuid
from typing import List, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.logger import logger as fastapi_logger

from fastmcp import MCPServer, MCPRequest, MCPResponse
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Qdrant
from langchain.schema import Document

from rich.console import Console
from rich.table import Table

# Настройки сервера
APP_NAME = "AgentMemoryMCP"
APP_VERSION = "1.0.0"
QDRANT_URL = "http://localhost:6333"
COLLECTION_NAME = "agent_memory"

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

# Создаём FastAPI приложение
app = FastAPI(title=APP_NAME, version=APP_VERSION)

# CORS (если понадобится)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Инициализация Qdrant
embeddings = OpenAIEmbeddings()
vectorstore = Qdrant(
    client=qdrant_client := Qdrant(
        url=QDRANT_URL,
        collection_name=COLLECTION_NAME,
        embeddings=embeddings,
    ),
    embeddings=embeddings,
    collection_name=COLLECTION_NAME,
)

# Создаём MCP сервер
mcp_server = MCPServer(app=app, name=APP_NAME, version=APP_VERSION)


# Вспомогательные функции
def _log_request(req: MCPRequest):
    table = Table(title="MCP Request", show_header=True, header_style="bold magenta")
    table.add_column("Method")
    table.add_column("Path")
    table.add_column("Body")
    table.add_row(req.method, req.path, json.dumps(req.body, indent=2))
    console.print(table)


def _log_response(res: MCPResponse):
    table = Table(title="MCP Response", show_header=True, header_style="bold green")
    table.add_column("Status")
    table.add_column("Body")
    table.add_row(str(res.status), json.dumps(res.body, indent=2))
    console.print(table)


# Обработчики MCP
@mcp_server.on("memory.add")
async def handle_memory_add(req: MCPRequest) -> MCPResponse:
    _log_request(req)
    content = req.body.get("content")
    if not content:
        raise HTTPException(status_code=400, detail="Missing 'content' field")
    # Создаём документ и сохраняем в Qdrant
    doc = Document(page_content=content, metadata={"id": str(uuid.uuid4())})
    vectorstore.add_documents([doc])
    res = MCPResponse(status=200, body={"id": doc.metadata["id"]})
    _log_response(res)
    return res


@mcp_server.on("memory.get")
async def handle_memory_get(req: MCPRequest) -> MCPResponse:
    _log_request(req)
    query = req.body.get("query")
    if not query:
        raise HTTPException(status_code=400, detail="Missing 'query' field")
    # Поиск похожих документов
    results = vectorstore.similarity_search(query, k=5)
    res_body = [
        {"id": doc.metadata.get("id"), "content": doc.page_content} for doc in results
    ]
    res = MCPResponse(status=200, body=res_body)
    _log_response(res)
    return res


@mcp_server.on("memory.delete")
async def handle_memory_delete(req: MCPRequest) -> MCPResponse:
    _log_request(req)
    doc_id = req.body.get("id")
    if not doc_id:
        raise HTTPException(status_code=400, detail="Missing 'id' field")
    # Удаляем документ по id
    deleted = vectorstore.delete(ids=[doc_id])
    res = MCPResponse(status=200, body={"deleted": deleted})
    _log_response(res)
    return res


# Запуск через uvicorn
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)