import os
import uuid
import datetime
from pymongo import MongoClient
from supabase import create_client, Client
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer

load_dotenv()

# MongoDB Configuration
MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017/")
MONGO_DB = "diet_db"

# Supabase Configuration
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY") or os.getenv("SUPABASE_KEY")

def migrate_kb():
    print("Starting Knowledge Base migration from MongoDB to Supabase...")
    
    if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
        print("Error: Supabase configuration missing.")
        return

    mongo_client = MongoClient(MONGO_URI)
    supabase: Client = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)
    
    # Load embedding model
    print("Loading embedding model (all-MiniLM-L6-v2)...")
    model = SentenceTransformer('all-MiniLM-L6-v2')

    # We need to migrate multiple collections possibly
    # In kb_manager.py, default was 'documents' in 'medical_kb'
    # In dietary_tools.py, it was 'diabetes' in 'test_index'
    collections_to_migrate = [
        {"db": "medical_kb", "coll": "documents"},
        {"db": "test_index", "coll": "diabetes"},
        {"db": "diet_db", "coll": "diabetes"}
    ]

    for item in collections_to_migrate:
        db_name = item["db"]
        coll_name = item["coll"]
        
        print(f"Migrating {db_name}.{coll_name}...")
        try:
            mongo_docs = list(mongo_client[db_name][coll_name].find())
            if not mongo_docs:
                print(f"No documents found in {db_name}.{coll_name}. Skipping.")
                continue

            # Batch processing for efficiency
            batch_size = 50
            for i in range(0, len(mongo_docs), batch_size):
                batch = mongo_docs[i:i+batch_size]
                texts = [doc.get("text", "") for doc in batch]
                
                # Generate embeddings
                embeddings = model.encode(texts).tolist()
                
                supabase_data = []
                for doc, emb in zip(batch, embeddings):
                    supabase_data.append({
                        "id": str(uuid.uuid4()),
                        "content": doc.get("text", ""),
                        "metadata": doc.get("metadata", {}),
                        "embedding": emb,
                        "collection_name": f"{db_name}_{coll_name}"
                    })
                
                supabase.table("knowledge_base").insert(supabase_data).execute()
                print(f"Uploaded {len(supabase_data)} documents...")

            print(f"Successfully migrated {len(mongo_docs)} documents from {db_name}.{coll_name}.")
        except Exception as e:
            print(f"Error migrating {db_name}.{coll_name}: {e}")

    print("Knowledge Base migration complete!")

if __name__ == "__main__":
    migrate_kb()