#!/usr/bin/env python3
"""
Database migration script to add practice_area fields to existing tables
Run with: /var/www/lawbot/venv/bin/python3 /var/www/lawbot/add_practice_area_fields.py
"""

import sys
import os
sys.path.insert(0, '/var/www/lawbot')

from app import app
from models import db
from sqlalchemy import text

def add_practice_area_columns():
    """Add practice_area columns to cases, analysis_jobs, and summarization_jobs tables"""

    with app.app_context():
        try:
            print("Starting database migration for practice area fields...")

            # Check if columns already exist before adding them
            inspector = db.inspect(db.engine)

            # Add practice_area to cases table
            cases_columns = [col['name'] for col in inspector.get_columns('cases')]
            if 'practice_area' not in cases_columns:
                print("Adding practice_area column to cases table...")
                with db.engine.connect() as conn:
                    conn.execute(text("""
                        ALTER TABLE cases
                        ADD COLUMN practice_area VARCHAR(50);
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_case_practice_area ON cases(practice_area);
                    """))
                    conn.commit()
                print("✅ Added practice_area to cases table")
            else:
                print("⚠️  practice_area column already exists in cases table")

            # Add practice_area and analysis_type to analysis_jobs table
            analysis_columns = [col['name'] for col in inspector.get_columns('analysis_jobs')]
            if 'practice_area' not in analysis_columns:
                print("Adding practice_area column to analysis_jobs table...")
                with db.engine.connect() as conn:
                    conn.execute(text("""
                        ALTER TABLE analysis_jobs
                        ADD COLUMN practice_area VARCHAR(50);
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_analysis_practice_area ON analysis_jobs(practice_area);
                    """))
                    conn.commit()
                print("✅ Added practice_area to analysis_jobs table")
            else:
                print("⚠️  practice_area column already exists in analysis_jobs table")

            if 'analysis_type' not in analysis_columns:
                print("Adding analysis_type column to analysis_jobs table...")
                with db.engine.connect() as conn:
                    conn.execute(text("""
                        ALTER TABLE analysis_jobs
                        ADD COLUMN analysis_type VARCHAR(100);
                    """))
                    conn.commit()
                print("✅ Added analysis_type to analysis_jobs table")
            else:
                print("⚠️  analysis_type column already exists in analysis_jobs table")

            # Add practice_area to summarization_jobs table
            summary_columns = [col['name'] for col in inspector.get_columns('summarization_jobs')]
            if 'practice_area' not in summary_columns:
                print("Adding practice_area column to summarization_jobs table...")
                with db.engine.connect() as conn:
                    conn.execute(text("""
                        ALTER TABLE summarization_jobs
                        ADD COLUMN practice_area VARCHAR(50);
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_summary_practice_area ON summarization_jobs(practice_area);
                    """))
                    conn.commit()
                print("✅ Added practice_area to summarization_jobs table")
            else:
                print("⚠️  practice_area column already exists in summarization_jobs table")

            print("\n✅ Database migration completed successfully!")
            print("\nNext steps:")
            print("1. Restart the application: systemctl restart epolaw")
            print("2. Test the new practice area features")

        except Exception as e:
            print(f"\n❌ Error during migration: {e}")
            print("\nTroubleshooting:")
            print("1. Check database connection")
            print("2. Ensure you have ALTER TABLE privileges")
            print("3. Review the error message above")
            return False

    return True

def verify_migration():
    """Verify that the migration was successful"""
    with app.app_context():
        try:
            # Test queries to ensure columns exist
            print("\nVerifying migration...")

            with db.engine.connect() as conn:
                # Test cases table
                result = conn.execute(text("SELECT practice_area FROM cases LIMIT 1"))
                print("✅ cases.practice_area verified")

                # Test analysis_jobs table
                result = conn.execute(text("SELECT practice_area, analysis_type FROM analysis_jobs LIMIT 1"))
                print("✅ analysis_jobs.practice_area and analysis_type verified")

                # Test summarization_jobs table
                result = conn.execute(text("SELECT practice_area FROM summarization_jobs LIMIT 1"))
                print("✅ summarization_jobs.practice_area verified")

            print("\n✅ All migrations verified successfully!")
            return True

        except Exception as e:
            print(f"\n❌ Verification failed: {e}")
            return False

if __name__ == "__main__":
    print("=" * 60)
    print("EPOLaw Practice Area Fields Migration")
    print("=" * 60)

    if add_practice_area_columns():
        verify_migration()

    print("\n" + "=" * 60)