import os
import psycopg2
from dotenv import load_dotenv

load_dotenv()

def migrate_db():
    try:
        conn = psycopg2.connect(
            host=os.getenv("DB_HOST"),
            port=os.getenv("DB_PORT"),
            database=os.getenv("DB_NAME"),
            user=os.getenv("DB_USER"),
            password=os.getenv("DB_PASSWORD")
        )
        with conn.cursor() as cur:
            # Add outlet_code column
            cur.execute("""
                ALTER TABLE google_review 
                ADD COLUMN IF NOT EXISTS outlet_code VARCHAR(255);
            """)
            print("Added 'outlet_code' column.")
            
            # Add foreign key constraint
            # Note: master_outlet(popcorn_code) must be unique. The inspection showed it has a unique constraint.
            try:
                cur.execute("""
                    ALTER TABLE google_review 
                    ADD CONSTRAINT fk_outlet_code 
                    FOREIGN KEY (outlet_code) 
                    REFERENCES master_outlet(popcorn_code) 
                    ON DELETE SET NULL;
                """)
                print("Added foreign key constraint for 'outlet_code'.")
            except psycopg2.errors.DuplicateObject:
                print("Foreign key constraint already exists.")
                conn.rollback()
            
            # Add language column
            cur.execute("""
                ALTER TABLE google_review 
                ADD COLUMN IF NOT EXISTS language VARCHAR(10);
            """)
            print("Added 'language' column.")
            
            conn.commit()
            print("Migration completed successfully.")
            
    except Exception as e:
        print(f"Error migrating database: {e}")
    finally:
        if conn: conn.close()

if __name__ == "__main__":
    migrate_db()