import pickle
import os

# Paths
CHUNKS_FILE = "sermon_chunks.pkl"
BACKUP_FILE = "sermon_chunks_backup.pkl"

def extract_date_code(filename: str) -> str:
    """Extracts date code from source, e.g., '55-0123A The Approach To God.pdf' -> '55-0123A'"""
    if not filename:
        return ""
    return filename.split()[0].replace(".pdf", "")

print("Loading chunks...")
with open(CHUNKS_FILE, "rb") as f:
    chunks = pickle.load(f)

# Save backup first
if not os.path.exists(BACKUP_FILE):
    print(f"Creating backup at {BACKUP_FILE}...")
    with open(BACKUP_FILE, "wb") as f:
        pickle.dump(chunks, f)

fixed_count = 0
for chunk in chunks:
    meta = chunk.metadata
    if not meta:
        continue
        
    source = meta.get("source", "")
    current_date_code = meta.get("date_code", "")
    current_chunk_id = meta.get("chunk_id", "")
    
    # Check if we can derive a better date_code from the source
    correct_date_code = extract_date_code(source)
    
    # If the current date code is wrong (e.g. '00-0000' or doesn't match the source exactly)
    if correct_date_code and current_date_code != correct_date_code:
        # We need to fix it!
        old_date_code = current_date_code
        meta["date_code"] = correct_date_code
        
        # We also need to fix the chunk_id. Usually formatted like: "00-0000|p74|part1-1|pg11-11"
        if current_chunk_id.startswith(old_date_code + "|"):
            meta["chunk_id"] = current_chunk_id.replace(old_date_code + "|", correct_date_code + "|", 1)
        elif old_date_code == "00-0000" and "00-0000|" in current_chunk_id:
            meta["chunk_id"] = current_chunk_id.replace("00-0000|", correct_date_code + "|", 1)
            
        fixed_count += 1

print(f"Identified and fixed {fixed_count} chunks with incorrect date_codes.")

if fixed_count > 0:
    print("Saving corrected chunks back to file...")
    with open(CHUNKS_FILE, "wb") as f:
        pickle.dump(chunks, f)
    print("Done! The .pkl file has been updated.")
else:
    print("No chunks needed fixing.")