import os
import chromadb
from sentence_transformers import SentenceTransformer
from loguru import logger

class SentenceTransformerEmbeddings:
    def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
        self.model = SentenceTransformer(model_name)

    def __call__(self, input: list[str]) -> list[list[float]]:
        embeddings = self.model.encode(input)
        return embeddings.tolist()

def load_documents():
    """Load and process documents into ChromaDB"""
    try:
        # Set up paths
        base_path = os.path.dirname(os.path.abspath(__file__))
        doc_path = os.path.join(base_path, 'a2023-45.txt')
        index_path = os.path.join(base_path, 'index.txt')
        chroma_path = os.path.join(base_path, 'chroma_db')

        # Ensure ChromaDB directory exists
        os.makedirs(chroma_path, exist_ok=True)
        
        logger.info(f"Loading documents from {doc_path} and {index_path}")

        # Initialize ChromaDB
        chroma_client = chromadb.PersistentClient(path=chroma_path)
        embedding_function = SentenceTransformerEmbeddings()
        
        # Create new collection (delete if exists)
        if "legal_documents" in [col.name for col in chroma_client.list_collections()]:
            chroma_client.delete_collection("legal_documents")
        
        collection = chroma_client.create_collection(
            name="legal_documents",
            embedding_function=embedding_function
        )

        # Read and validate files
        with open(doc_path, 'r', encoding='utf-8') as f:
            document = f.read().strip()
        
        with open(index_path, 'r', encoding='utf-8') as f:
            index_content = [line.strip() for line in f.readlines() if line.strip()]

        # Process document into sections
        sections = []
        current_section = ""
        current_title = ""
        
        for line in document.split('\n'):
            line = line.strip()
            if any(index_line in line for index_line in index_content):
                if current_section and current_title:
                    sections.append({
                        "title": current_title,
                        "content": current_section.strip()
                    })
                current_title = line
                current_section = ""
            else:
                if line:
                    current_section += line + "\n"
        
        # Add final section
        if current_section and current_title:
            sections.append({
                "title": current_title,
                "content": current_section.strip()
            })

        # Prepare data for ChromaDB
        documents = []
        metadatas = []
        ids = []
        
        for i, section in enumerate(sections):
            if section["content"].strip():
                documents.append(section["content"])
                metadatas.append({
                    "title": section["title"],
                    "source": "a2023-45.txt",
                    "section_number": i + 1
                })
                ids.append(f"section_{i+1}")

        # Add to ChromaDB
        collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )
        
        logger.info(f"Successfully loaded {len(documents)} sections into ChromaDB")
        return True

    except Exception as e:
        logger.error(f"Error loading documents: {str(e)}")
        return False

if __name__ == "__main__":
    success = load_documents()
    if success:
        print("Documents successfully loaded into ChromaDB")
    else:
        print("Failed to load documents into ChromaDB")