import os
import time
import sys
import json
import dotenv

from pinecone import Pinecone
from app import load_chunks, get_vector_store

# Load environment variables
dotenv.load_dotenv()

CHECKPOINT_FILE = "upload_checkpoint.json"

BATCH_SIZE = 64  # Optimally sized for high-precision paragraph chunks (averaging ~230 tokens each)
SLEEP_BETWEEN_BATCHES = 0.2  # Throttling is now managed by the high-speed Google Gemini API
MAX_RETRIES = 5
BACKOFF_FACTOR = 2

def clear_index():
    api_key = os.getenv("PINECONE_API_KEY")
    if not api_key:
        print("Error: PINECONE_API_KEY is not set.")
        sys.exit(1)
        
    pc = Pinecone(api_key=api_key)
    index_name = "branham-index"
    
    print(f"Connecting to Pinecone index '{index_name}'...")
    idx = pc.Index(index_name)
    
    stats = idx.describe_index_stats()
    total_vectors = stats.get("total_vector_count", 0)
    print(f"Current total vector count before clear: {total_vectors}")
    
    print(f"Deleting all vectors in '{index_name}' to start fresh...")
    try:
        idx.delete(delete_all=True)
    except Exception as e:
        print(f"Index deletion skipped (it might already be completely empty): {e}")
    
    # Wait for deletion to reflect
    time.sleep(5)
    stats = idx.describe_index_stats()
    print(f"Post-clear total vector count: {stats.get('total_vector_count', 0)}")
    print("Pinecone index cleared successfully!\n")

def upload_in_batches(chunks, start_idx=0):
    # Model is loaded inside get_vector_store(), track the exact loading time
    print("Initializing embedding model and Pinecone connection...")
    model_load_start = time.perf_counter()
    vector_store = get_vector_store()
    print(f"Initialization complete in {time.perf_counter() - model_load_start:.1f}s.\n")
    
    total_chunks = len(chunks)
    
    print(f"Starting batch upload of {total_chunks} chunks to Pinecone...")
    print(f"Batch Size: {BATCH_SIZE} | Spacer delay: {SLEEP_BETWEEN_BATCHES}s | Max Retries: {MAX_RETRIES}\n")
    
    # Track elapsed time strictly for encoding + uploading
    start_time = time.perf_counter()
    
    for i in range(start_idx, total_chunks, BATCH_SIZE):
        batch = chunks[i : i + BATCH_SIZE]
        batch_num = (i // BATCH_SIZE) + 1
        total_batches = (total_chunks + BATCH_SIZE - 1) // BATCH_SIZE
        
        # Implement robust retry with exponential backoff
        retries = 0
        while retries <= MAX_RETRIES:
            try:
                # Add to Pinecone
                vector_store.add_documents(batch)
                
                # Progress logging
                elapsed = time.perf_counter() - start_time
                pct = (min(i + BATCH_SIZE, total_chunks) / total_chunks) * 100
                rate = (i + len(batch) - start_idx) / elapsed if elapsed > 0 else 0
                eta = (total_chunks - (i + len(batch))) / rate if rate > 0 else 0
                
                print(
                    f"[{pct:6.2f}%] Batch {batch_num}/{total_batches} uploaded successfully. "
                    f"({min(i + BATCH_SIZE, total_chunks)}/{total_chunks}) | "
                    f"Speed: {rate:.1f} chunks/sec | ETA: {eta/60:.1f} min"
                )
                
                # Write checkpoint to file
                with open(CHECKPOINT_FILE, "w", encoding="utf-8") as f:
                    json.dump({"last_uploaded_index": i + len(batch)}, f)
                break
            except Exception as e:
                retries += 1
                if retries > MAX_RETRIES:
                    print(f"\n[FATAL ERROR] Batch {batch_num} failed completely after {MAX_RETRIES} retries. Error: {e}")
                    raise e
                    
                sleep_time = BACKOFF_FACTOR ** retries
                print(
                    f"\n[WARNING] Error on Batch {batch_num} (Attempt {retries}/{MAX_RETRIES}): {e}. "
                    f"Retrying in {sleep_time}s..."
                )
                time.sleep(sleep_time)
                
        time.sleep(SLEEP_BETWEEN_BATCHES)
        
    total_time = time.perf_counter() - start_time
    print(f"\nSUCCESS! Uploaded {total_chunks} chunks to Pinecone in {total_time/60:.2f} minutes.")
    
    # Delete checkpoint on success
    if os.path.exists(CHECKPOINT_FILE):
        os.remove(CHECKPOINT_FILE)

def main():
    # 1. Load the clean serialized chunks
    print("Loading chunks from 'sermon_chunks.pkl'...")
    chunks = load_chunks()
    if not chunks:
        print("Error: No chunks found in 'sermon_chunks.pkl'. Make sure chunking is complete.")
        return
        
    print(f"Successfully loaded {len(chunks)} sermon chunks.\n")
    
    # 2. Check for checkpoint to see if we can resume
    start_idx = 0
    if os.path.exists(CHECKPOINT_FILE):
        try:
            with open(CHECKPOINT_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
                start_idx = data.get("last_uploaded_index", 0)
        except Exception as e:
            print(f"[WARNING] Could not read checkpoint file: {e}. Starting fresh.")
            
    if start_idx > 0 and start_idx < len(chunks):
        print(f"[RESUME] Checkpoint found! Resuming upload from chunk index {start_idx}...")
    else:
        # 3. Clear index only when starting fresh
        print("[START FRESH] Starting fresh. Clearing index first...")
        clear_index()
        
    # 4. Upload chunks
    upload_in_batches(chunks, start_idx=start_idx)

if __name__ == "__main__":
    main()