#!/usr/bin/env python3
"""
Migration script to add email verification and location fields to database
Run this script to update your existing database with new fields
"""
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import create_engine, text
from config.database import DATABASE_URL
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def run_migration():
    """Add new fields to existing database"""
    engine = create_engine(DATABASE_URL)
    
    with engine.connect() as conn:
        # Start transaction
        trans = conn.begin()
        
        try:
            # Check if we're using PostgreSQL or SQLite
            is_postgres = 'postgresql' in DATABASE_URL
            
            # Add fields to users table
            logger.info("Adding new fields to users table...")
            
            # Check if columns already exist before adding
            if is_postgres:
                result = conn.execute(text("""
                    SELECT column_name 
                    FROM information_schema.columns 
                    WHERE table_name = 'users'
                """))
                existing_columns = [row[0] for row in result]
            else:
                result = conn.execute(text("PRAGMA table_info(users)"))
                existing_columns = [row[1] for row in result]
            
            if 'phone_number' not in existing_columns:
                conn.execute(text("ALTER TABLE users ADD COLUMN phone_number VARCHAR"))
                logger.info("Added phone_number column")
            
            if 'email_verification_token' not in existing_columns:
                conn.execute(text("ALTER TABLE users ADD COLUMN email_verification_token VARCHAR"))
                logger.info("Added email_verification_token column")
            
            if 'email_verified_at' not in existing_columns:
                if is_postgres:
                    conn.execute(text("ALTER TABLE users ADD COLUMN email_verified_at TIMESTAMP"))
                else:
                    conn.execute(text("ALTER TABLE users ADD COLUMN email_verified_at DATETIME"))
                logger.info("Added email_verified_at column")
            
            if 'verification_sent_at' not in existing_columns:
                if is_postgres:
                    conn.execute(text("ALTER TABLE users ADD COLUMN verification_sent_at TIMESTAMP"))
                else:
                    conn.execute(text("ALTER TABLE users ADD COLUMN verification_sent_at DATETIME"))
                logger.info("Added verification_sent_at column")
            
            # Add fields to patients table
            logger.info("Adding location fields to patients table...")
            
            if is_postgres:
                result = conn.execute(text("""
                    SELECT column_name 
                    FROM information_schema.columns 
                    WHERE table_name = 'patients'
                """))
                existing_columns = [row[0] for row in result]
            else:
                result = conn.execute(text("PRAGMA table_info(patients)"))
                existing_columns = [row[1] for row in result]
            
            location_fields = [
                ('address_line1', 'VARCHAR'),
                ('address_line2', 'VARCHAR'),
                ('city', 'VARCHAR'),
                ('state_province', 'VARCHAR'),
                ('postal_code', 'VARCHAR'),
                ('country', "VARCHAR DEFAULT 'USA'"),
                ('timezone', "VARCHAR DEFAULT 'America/New_York'"),
                ('medical_history', "TEXT DEFAULT ''")
            ]
            
            for field_name, field_type in location_fields:
                if field_name not in existing_columns:
                    conn.execute(text(f"ALTER TABLE patients ADD COLUMN {field_name} {field_type}"))
                    logger.info(f"Added {field_name} column")
            
            # Update therapy_goals from JSON to TEXT if needed
            if 'therapy_goals' in existing_columns:
                logger.info("Note: therapy_goals column exists. Consider manual data migration if changing from JSON to TEXT")
            
            trans.commit()
            logger.info("Migration completed successfully!")
            
        except Exception as e:
            trans.rollback()
            logger.error(f"Migration failed: {str(e)}")
            raise

if __name__ == "__main__":
    run_migration()