import os
import sys
import psycopg2
from pathlib import Path

def get_db_url():
    """Get database URL from environment variables."""
    db_user = os.getenv('DB_USER', 'mold_user')
    db_password = os.getenv('DB_PASSWORD', 'your_db_password_here')
    db_host = os.getenv('DB_HOST', 'localhost')
    db_port = os.getenv('DB_PORT', '5434')
    db_name = os.getenv('DB_NAME', 'mold_cost_calculator_dev')
    
    return f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"

def test_migration():
    # Get database URL from environment
    db_url = get_db_url()
    
    try:
        # Connect to the database
        conn = psycopg2.connect(db_url)
        cur = conn.cursor()
        
        # Check if the column exists
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'mold_quotations' 
            AND column_name = 'mold_specific_type'
        """)
        
        if cur.fetchone():
            print("✅ mold_specific_type column exists")
        else:
            print("❌ mold_specific_type column does not exist")
            return False
        
        # Check if the column is NOT NULL
        cur.execute("""
            SELECT is_nullable 
            FROM information_schema.columns 
            WHERE table_name = 'mold_quotations' 
            AND column_name = 'mold_specific_type'
        """)
        
        is_nullable = cur.fetchone()[0]
        if is_nullable == 'NO':
            print("✅ mold_specific_type column is NOT NULL")
        else:
            print("❌ mold_specific_type column is nullable")
            return False
        
        # Check if all records have a value
        cur.execute("""
            SELECT COUNT(*) 
            FROM mold_quotations 
            WHERE mold_specific_type IS NULL
        """)
        
        null_count = cur.fetchone()[0]
        if null_count == 0:
            print("✅ All records have a mold_specific_type value")
        else:
            print(f"❌ Found {null_count} records with NULL mold_specific_type")
            return False
        
        print("All tests passed!")
        return True
        
    except Exception as e:
        print(f"Error testing migration: {str(e)}")
        return False
    finally:
        if 'cur' in locals():
            cur.close()
        if 'conn' in locals():
            conn.close()

if __name__ == '__main__':
    success = test_migration()
    sys.exit(0 if success else 1) 