from fastmcp import FastMCP
import chromadb
from chromadb.utils import embedding_functions
from datetime import datetime
import os
import shutil
from huggingface_hub import HfApi, snapshot_download

# --- КОНФИГУРАЦИЯ ---
# Локальная временная папка (внутри контейнера)
LOCAL_DB_PATH = "/app/chroma_db_cache"

# Твой личный сейф для памяти
REPO_ID = "LevinAleksey/neuron-brain-storage" 

# Токен берется из Secrets спейса (убедись, что добавил HF_TOKEN)
HF_TOKEN = os.environ.get("HF_TOKEN")

# --- ЛОГИКА СИНХРОНИЗАЦИИ (ОБЛАКО <-> СЕРВЕР) ---
api = HfApi(token=HF_TOKEN)

def pull_memory():
    """Скачиваем память из Dataset при старте сервера"""
    print(f"📥 PULL: Подключаюсь к хранилищу {REPO_ID}...")
    try:
        # Проверка: если репозиторий пустой, не пытаемся качать
        if not api.repo_info(repo_id=REPO_ID, repo_type="dataset").sha:
            print("⚠️ Репозиторий памяти пуст. Начинаем с чистого листа.")
            return

        # Чистим локальный кэш перед загрузкой во избежание конфликтов
        if os.path.exists(LOCAL_DB_PATH):
            shutil.rmtree(LOCAL_DB_PATH)
            
        # Скачиваем файлы базы данных
        snapshot_download(
            repo_id=REPO_ID, 
            repo_type="dataset", 
            local_dir=LOCAL_DB_PATH,
            token=HF_TOKEN,
            allow_patterns=["chroma*"] # Качаем только файлы базы
        )
        print("✅ PULL: Память успешно восстановлена!")
    except Exception as e:
        print(f"ℹ️ Инфо (первый запуск или ошибка): {e}")

def push_memory(reason="Auto-save"):
    """Отправляем обновления памяти обратно в Dataset"""
    print(f"☁️ PUSH: Синхронизация с облаком ({reason})...")
    try:
        api.upload_folder(
            folder_path=LOCAL_DB_PATH,
            repo_id=REPO_ID,
            repo_type="dataset",
            path_in_repo=".",
            commit_message=f"Memory Sync: {reason}",
            token=HF_TOKEN
        )
        print("✅ PUSH: Знания надежно сохранены в Dataset!")
    except Exception as e:
        print(f"❌ Ошибка сохранения в облако: {e}")

# 1. СТАРТ: Восстанавливаем память
pull_memory()

# 2. ИНИЦИАЛИЗАЦИЯ: Подключаем ChromaDB
chroma_client = chromadb.PersistentClient(path=LOCAL_DB_PATH)
emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)
collection = chroma_client.get_or_create_collection(
    name="architect_memory",
    embedding_function=emb_fn
)

# 3. СЕРВЕР MCP
mcp = FastMCP("Neuron_Architect_Brain")

# --- ИНСТРУМЕНТЫ (TOOLS) ---

@mcp.tool()
def learn_update(project: str, summary: str, details: str = ""):
    """
    [ВАЖНО] Используй этот инструмент, чтобы ЗАПОМНИТЬ архитектурные изменения.
    project: имя проекта (School, Clinic, etc.)
    summary: суть изменений
    details: технические детали или код
    """
    timestamp = datetime.now().isoformat()
    text_chunk = f"UPDATE: {timestamp}\nPROJECT: {project}\nSUMMARY: {summary}\nDETAILS: {details}"
    
    # 1. Сохраняем локально
    collection.add(
        documents=[text_chunk],
        metadatas=[{"project": project, "type": "update", "timestamp": timestamp}],
        ids=[f"{project}_{datetime.now().timestamp()}"]
    )
    
    # 2. Сохраняем в облако (Dataset)
    push_memory(f"Update for {project}")
    
    return f"✅ Запомнил и синхронизировал с репозиторием {REPO_ID}"

@mcp.tool()
def recall_context(query: str, project: str = None):
    """
    Используй этот инструмент, чтобы ВСПОМНИТЬ прошлые решения.
    """
    where_filter = {"project": project} if project else None
    
    results = collection.query(
        query_texts=[query], 
        n_results=3, 
        where=where_filter
    )
    
    if not results['documents'] or not results['documents'][0]:
        return "В памяти ничего не найдено по этому запросу."
        
    return "\n---\n".join(results['documents'][0])

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