import os
import sys
import uuid
import datetime
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
from supabase import create_client, Client
from src.utils.logger import setup_logger

logger = setup_logger("KBManager")

def _get_supabase_client() -> Client:
    url = os.getenv("SUPABASE_URL")
    key = os.getenv("SUPABASE_SERVICE_ROLE_KEY") or os.getenv("SUPABASE_KEY")
    if not url or not key:
        raise ValueError("Supabase configuration missing")
    return create_client(url, key)

def create_knowledge_base_supabase(pdf_path: str, collection_name: str = "general"):
    """
    Creates a new knowledge base in Supabase (pgvector) from a PDF document.
    """
    if not os.path.exists(pdf_path):
        logger.error(f"Error: File not found at {pdf_path}")
        return
        
    logger.info(f"Loading document: {pdf_path}")
    loader = PyPDFLoader(pdf_path)
    docs = loader.load()
    
    for doc in docs:
        if 'page' in doc.metadata:
            doc.metadata['page_index'] = doc.metadata['page']
            
    logger.info(f"Successfully loaded {len(docs)} pages.")
    
    logger.info("Splitting text into chunks...")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)
    logger.info(f"Created {len(splits)} chunks.")
    
    # Generate embeddings
    logger.info("Generating embeddings...")
    model = SentenceTransformer('all-MiniLM-L6-v2')
    texts = [s.page_content for s in splits]
    embeddings = model.encode(texts).tolist()
    
    supabase = _get_supabase_client()
    
    # Store in Supabase
    logger.info(f"Storing {len(splits)} chunks in Supabase knowledge_base...")
    supabase_data = []
    for chunk, emb in zip(splits, embeddings):
        supabase_data.append({
            "content": chunk.page_content,
            "metadata": chunk.metadata,
            "embedding": emb,
            "collection_name": collection_name
        })
        
    # Bulk insert (limited by payload size, so batching)
    batch_size = 100
    for i in range(0, len(supabase_data), batch_size):
        supabase.table("knowledge_base").insert(supabase_data[i:i+batch_size]).execute()
        
    logger.info("Knowledge base integration to Supabase complete.")
    return True

if __name__ == "__main__":
    # Example usage for Supabase migration
    sources_dir = "data/sources"
    if os.path.exists(sources_dir):
        for filename in os.listdir(sources_dir):
            if filename.endswith(".pdf"):
                pdf_path = os.path.join(sources_dir, filename)
                create_knowledge_base_supabase(pdf_path, collection_name="diabetes_guidelines")