import sys
sys.path.append('/app')
from sqlmodel import Session, select
from database import engine
import models

def fix_text(text):
    if not text: return text
    
    # Fix NO3- where 3 is on the next line
    text = text.replace('NO−\n3', '$NO_3^-$')
    text = text.replace('NO-\n3', '$NO_3^-$')
    
    text = text.replace('CO2−\n3', '$CO_3^{2-}$')
    text = text.replace('CO2-\n3', '$CO_3^{2-}$')
    
    text = text.replace('SO2−\n4', '$SO_4^{2-}$')
    text = text.replace('SO2-\n4', '$SO_4^{2-}$')
    
    text = text.replace('PO3−\n4', '$PO_4^{3-}$')
    text = text.replace('PO3-\n4', '$PO_4^{3-}$')
    
    # Also generic charges if they are on same line
    text = text.replace('NO3-', '$NO_3^-$')
    text = text.replace('CO32-', '$CO_3^{2-}$')
    text = text.replace('SO42-', '$SO_4^{2-}$')
    text = text.replace('PO43-', '$PO_4^{3-}$')
    text = text.replace('O2-', '$O^{2-}$')
    text = text.replace('F-', '$F^-$')
    text = text.replace('Cl-', '$Cl^-$')
    text = text.replace('Br-', '$Br^-$')
    text = text.replace('I-', '$I^-$')
    text = text.replace('Na+', '$Na^+$')
    text = text.replace('K+', '$K^+$')
    text = text.replace('Ca2+', '$Ca^{2+}$')
    text = text.replace('Mg2+', '$Mg^{2+}$')
    text = text.replace('Fe2+', '$Fe^{2+}$')
    text = text.replace('Fe3+', '$Fe^{3+}$')
    text = text.replace('Cu2+', '$Cu^{2+}$')
    text = text.replace('Al3+', '$Al^{3+}$')
    
    # Check if we have O3 (ozone) that wasn't replaced
    text = text.replace('O3', '$O_3$')

    return text

with Session(engine) as session:
    questions = session.exec(select(models.Question)).all()
    count = 0
    for q in questions:
        old_text = q.question_text
        new_text = fix_text(old_text)
        if old_text != new_text:
            q.question_text = new_text
            session.add(q)
            count += 1
            print(f'Updated Q {q.id}')
    session.commit()
    print(f'Successfully updated {count} questions for ions.')