from sqlalchemy import create_engine, text
from app.core.config import settings

def update_schema():
    engine = create_engine(settings.DATABASE_URL)
    with engine.connect() as conn:
        print("Adding token tracking columns to modules table...")
        try:
            conn.execute(text("ALTER TABLE modules ADD COLUMN ingress_tokens INTEGER DEFAULT 0"))
            print("Added ingress_tokens column.")
        except Exception as e:
            print(f"ingress_tokens column might already exist: {e}")

        try:
            conn.execute(text("ALTER TABLE modules ADD COLUMN egress_tokens INTEGER DEFAULT 0"))
            print("Added egress_tokens column.")
        except Exception as e:
            print(f"egress_tokens column might already exist: {e}")

        try:
            conn.execute(text("ALTER TABLE modules ADD COLUMN total_tokens INTEGER DEFAULT 0"))
            print("Added total_tokens column.")
        except Exception as e:
            print(f"total_tokens column might already exist: {e}")
            
        conn.commit()
        print("Schema update complete.")

if __name__ == "__main__":
    update_schema()