import os
import pymysql
from database import DB_CONFIG

def migrate_db():
    conn = pymysql.connect(
        host=DB_CONFIG['host'],
        port=int(DB_CONFIG['port']),
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        database=DB_CONFIG['database']
    )
    
    try:
        with conn.cursor() as cursor:
            # Check if column exists
            print("Checking schema...")
            cursor.execute("SHOW COLUMNS FROM ktp_records LIKE 'image_path'")
            result = cursor.fetchone()
            
            if not result:
                print("Adding image_path column to ktp_records...")
                cursor.execute("ALTER TABLE ktp_records ADD COLUMN image_path VARCHAR(255) NULL AFTER berlaku_hingga")
                conn.commit()
                print("Migration successful: Added image_path column.")
            else:
                print("Column image_path already exists in KTP. No migration needed.")
            
            # Check KK
            print("Checking KK schema...")
            cursor.execute("SHOW COLUMNS FROM kk_records LIKE 'image_path'")
            result_kk = cursor.fetchone()
            
            if not result_kk:
                print("Adding image_path column to kk_records...")
                cursor.execute("ALTER TABLE kk_records ADD COLUMN image_path VARCHAR(255) NULL AFTER kode_pos")
                conn.commit()
                print("Migration successful: Added image_path column to KK.")
            else:
                print("Column image_path already exists in KK.")
                
    except Exception as e:
        print(f"Migration error: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    migrate_db()