import os
import chromadb
from db import SanatanDatabase
from config import SanatanConfig
import logging
import argparse

# Configure logging
logging.basicConfig(level=logging.INFO, force=True)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

def get_target_field(scripture_config):
    """Identifies the hierarchical index field from scripture config."""
    allowed_indices = ["relative_verse_number", "sloka_number", "sloka_index", "verse_number", "position_in_chapter"]
    for field in scripture_config.get("metadata_fields", []):
        if field["name"] in allowed_indices:
            return field["name"]
    return None

def migrate_position_in_chapter(full_run=False):
    config = SanatanConfig()
    db = SanatanDatabase()
    
    # Use HttpClient with CHROMADB_SPACE_URL
    chroma_url = os.getenv("CHROMADB_SPACE_URL")
    if not chroma_url:
        logger.error("CHROMADB_SPACE_URL not set in environment.")
        return
        
    chroma_client = chromadb.HttpClient(host=chroma_url)
    
    BATCH_SIZE = 500
    
    for scripture in config.scriptures:
        scripture_name = scripture["name"]
        collection_name = scripture["collection_name"]
        target_field = get_target_field(scripture)

        logger.info(f"Processing {scripture_name} (Chroma: {collection_name}, Field: {target_field or '_global_index'})")
        
        # If not a full run, check if any verses already have the property
        if not full_run:
            count_query = f"""
            MATCH (v:Verse)-[:PART_OF]->(s:Scripture {{name: '{scripture_name}'}})
            WHERE v.position_in_chapter IS NOT NULL
            RETURN count(v) as count
            """
            result = db.run_cypher(count_query)
            existing_count = result[0]["count"] if result and "count" in result[0] else 0
            logger.info(f"existing_count = {existing_count}")
            if existing_count > 0:
                logger.info(f"Skipping {scripture_name}: already has {existing_count} verses populated. Use --full to override.")
                continue

        # Check if collection exists exactly as named
        coll_name = scripture["collection_name"]
        try:
            collection = chroma_client.get_collection(coll_name)
        except:
            # Fallback for _openai suffix
            try:
                coll_name = f"{coll_name}_openai"
                collection = chroma_client.get_collection(coll_name)
            except:
                logger.warning(f"Could not find collection for {scripture_name}")
                continue
            
        logger.info(f"Processing {scripture_name} (Chroma: {coll_name}, Field: {target_field or '_global_index'})")
        
        # Use a loop to fetch all records in case of pagination limits
        data = {"metadatas": [], "ids": []}
        offset = 0
        limit = 1000
        while True:
            chunk = collection.get(offset=offset, limit=limit, include=["metadatas"])
            if not chunk["ids"]:
                break
            data["metadatas"].extend(chunk["metadatas"])
            data["ids"].extend(chunk["ids"])
            offset += limit
            
        metadatas = data["metadatas"]
        
        logger.info(f"Retrieved {len(metadatas)} records for {scripture_name}")
        
        updates = []
        for i, meta in enumerate(metadatas):
            # The global_id in ArcadeDB is consistently scripture_name + "_" + _global_index
            global_index = meta.get("_global_index")
            if global_index is None:
                continue
            
            global_id = f"{scripture_name}_{global_index}"
            
            # Use dictionary get safely for relative_verse_number etc
            val = meta.get(target_field) if target_field and target_field in meta else global_index
            
            if i == 0:
                logger.info(f"DEBUG: scripture={scripture_name}, global_id={global_id}, target={target_field}, val={val}")
                
            if val is not None:
                updates.append({"global_id": global_id, "position": val})
        
        if updates:
            # DEBUG: Sample the first update to ensure it's not None
            logger.info(f"DEBUG: Sample update: {updates[0]}")
            logger.info(f"Updating {len(updates)} records for {scripture_name} in batches of {BATCH_SIZE}...")
            
            for i in range(0, len(updates), BATCH_SIZE):
                batch = updates[i:i + BATCH_SIZE]
                
                query = """
                UNWIND $batch AS row
                MATCH (v:Verse {global_id: row.global_id})
                SET v.position_in_chapter = row.position
                """
                try:
                    db.run_cypher(query, {"batch": batch})
                    logger.info(f"Batch {i//BATCH_SIZE + 1} for {scripture_name} successful.")
                except Exception as e:
                    logger.error(f"Failed to update batch starting at {i}: {e}")
        else:
            logger.info(f"No updates found for {scripture_name}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Migrate position_in_chapter to ArcadeDB")
    parser.add_argument("--full", action="store_true", help="Perform a full run, overwriting existing data")
    args = parser.parse_args()
    
    migrate_position_in_chapter(full_run=args.full)