from sqlalchemy import create_engine, text
from app.config import DATABASE_URL, TABLE_PREFIX

def update_schema():
    engine = create_engine(DATABASE_URL)
    with engine.connect() as conn:
        print("Updating schema...")
        
        # Add erp_account to DitRecord
        try:
            conn.execute(text(f"ALTER TABLE {TABLE_PREFIX}DIT_Records ADD COLUMN erp_account VARCHAR(100)"))
            print("Added erp_account to DIT_Records")
        except Exception as e:
            print(f"erp_account might already exist: {e}")

        # Add oppy_no and cust_id to SampleRecord
        try:
            conn.execute(text(f"ALTER TABLE {TABLE_PREFIX}Sample_Records ADD COLUMN oppy_no VARCHAR(100)"))
            print("Added oppy_no to Sample_Records")
        except Exception as e:
            print(f"oppy_no might already exist: {e}")
            
        try:
            conn.execute(text(f"ALTER TABLE {TABLE_PREFIX}Sample_Records ADD COLUMN cust_id VARCHAR(100)"))
            print("Added cust_id to Sample_Records")
        except Exception as e:
            print(f"cust_id might already exist: {e}")

        # Add cust_id to OrderRecord
        try:
            conn.execute(text(f"ALTER TABLE {TABLE_PREFIX}Order_Records ADD COLUMN cust_id VARCHAR(100)"))
            print("Added cust_id to Order_Records")
        except Exception as e:
            print(f"cust_id might already exist: {e}")

        # Add match_priority and match_source to MatchResult
        try:
            conn.execute(text(f"ALTER TABLE {TABLE_PREFIX}Match_Results ADD COLUMN match_priority INTEGER DEFAULT 3"))
            print("Added match_priority to Match_Results")
        except Exception as e:
            print(f"match_priority might already exist: {e}")
            
        try:
            conn.execute(text(f"ALTER TABLE {TABLE_PREFIX}Match_Results ADD COLUMN match_source VARCHAR(255)"))
            print("Added match_source to Match_Results")
        except Exception as e:
            print(f"match_source might already exist: {e}")

        conn.commit()
        print("Schema update completed.")

if __name__ == "__main__":
    update_schema()