import json
import os
import argparse
import chromadb
from chromadb.utils import embedding_functions
from typing import List, Dict
from mcp.server.fastmcp import FastMCP

# Inicializa o servidor FastMCP
mcp = FastMCP("Fluig Technical Wiki")

# Configurações do Banco Vetorial
DB_PATH = "fluig_vector_db"
SNIPPETS_DIR = os.path.join("fluig_rag_docs", "Biblioteca de Snippets")

# Inicialização preguiçosa (lazy load) do ChromaDB para economizar memória no Pi
_collection = None

def get_collection():
    global _collection
    if _collection is None:
        if os.path.exists(DB_PATH):
            model_name = "all-MiniLM-L6-v2"
            emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)
            client = chromadb.PersistentClient(path=DB_PATH)
            _collection = client.get_collection(name="fluig_docs", embedding_function=emb_fn)
        else:
            print("AVISO: Banco vetorial não encontrado. Use --mode stdio para busca simples ou gere o banco.")
    return _collection

@mcp.tool()
def search_docs(query: str) -> str:
    """
    Busca semântica profunda na documentação técnica do Fluig usando Vetores. 
    Encontra resultados por significado, mesmo que as palavras exatas não coincidam.
    """
    collection = get_collection()
    
    if collection:
        # Busca Semântica
        results = collection.query(
            query_texts=[query],
            n_results=5
        )
        
        output = []
        for i in range(len(results['documents'][0])):
            doc = results['documents'][0][i]
            meta = results['metadatas'][0][i]
            output.append(f"### {meta.get('title')}\n{doc}\n---\n")
        return "\n".join(output)
    else:
        # Fallback para busca por palavra-chave se o banco vetorial não estiver pronto
        return "Banco vetorial não inicializado. Por favor, execute vector_db_manager.py no servidor."

@mcp.tool()
def get_code_snippets(language: str) -> str:
    """Recupera exemplos de código Fluig (javascript, java ou sql)."""
    lang = language.lower()
    file_map = {"javascript": "Snippets JAVASCRIPT.md", "js": "Snippets JAVASCRIPT.md", "java": "Snippets JAVA.md", "sql": "Snippets SQL.md"}
    file_name = file_map.get(lang)
    if not file_name: return f"Linguagem '{language}' não suportada."
    path = os.path.join(SNIPPETS_DIR, file_name)
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f: return f.read()
    return "Snippets não encontrados."

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fluig MCP Server with Vector Search")
    parser.add_argument("--mode", choices=["stdio", "sse"], default="stdio")
    parser.add_argument("--port", type=int, default=8001)
    args = parser.parse_args()
    
    if args.mode == "sse":
        mcp.run(transport="sse")
    else:
        mcp.run(transport="stdio")