#!/usr/bin/env python3
import sqlite3
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from models import db, User, Company, AnalysisJob, AnalysisResult
from datetime import datetime

# Create MySQL engine
mysql_engine = create_engine('mysql+pymysql://lawbot_user:AK$82Z5FOn38i&!@localhost/lawbot')

# Create all tables in MySQL
db.metadata.create_all(mysql_engine)

# Connect to SQLite
sqlite_conn = sqlite3.connect('lawbot.db')
sqlite_conn.row_factory = sqlite3.Row
sqlite_cursor = sqlite_conn.cursor()

# Create MySQL session
Session = sessionmaker(bind=mysql_engine)
mysql_session = Session()

print("Starting migration from SQLite to MySQL...")

try:
    # Check if company table exists in SQLite
    sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='company'")
    if sqlite_cursor.fetchone():
        print("Migrating companies...")
        sqlite_cursor.execute("SELECT * FROM company")
        companies = sqlite_cursor.fetchall()
        for row in companies:
            # Get column names
            columns = [column[0] for column in sqlite_cursor.description]
            row_dict = dict(zip(columns, row))
            
            company = Company(
                id=row_dict['id'],
                name=row_dict['name'],
                description=row_dict.get('description', ''),
                subscription_tier='free',
                created_at=datetime.utcnow()
            )
            mysql_session.merge(company)
        mysql_session.commit()
        print(f"Migrated {len(companies)} companies")

    # Migrate Users
    print("Migrating users...")
    sqlite_cursor.execute("SELECT * FROM user")
    users = sqlite_cursor.fetchall()
    columns = [column[0] for column in sqlite_cursor.description]
    
    for row in users:
        row_dict = dict(zip(columns, row))
        user = User(
            id=row_dict['id'],
            username=row_dict['username'],
            email=row_dict['email'],
            password_hash=row_dict['password_hash'],
            role=row_dict.get('role', 'user'),
            company_id=row_dict.get('company_id'),
            created_at=datetime.utcnow()
        )
        mysql_session.merge(user)
    mysql_session.commit()
    print(f"Migrated {len(users)} users")

    # Migrate AnalysisJobs
    print("Migrating analysis jobs...")
    sqlite_cursor.execute("SELECT * FROM analysis_job")
    jobs = sqlite_cursor.fetchall()
    if jobs:
        columns = [column[0] for column in sqlite_cursor.description]
        
        for row in jobs:
            row_dict = dict(zip(columns, row))
            
            # Get company_id from user if not in job
            company_id = row_dict.get('company_id')
            if not company_id and row_dict.get('user_id'):
                user = mysql_session.query(User).filter_by(id=row_dict['user_id']).first()
                if user:
                    company_id = user.company_id
            
            job = AnalysisJob(
                job_uuid=row_dict['id'],  # Using existing ID as UUID
                user_id=row_dict['user_id'],
                company_id=company_id,
                filename=row_dict['filename'],
                original_filename=row_dict.get('filename', ''),
                perspective=row_dict['perspective'],
                status=row_dict['status'],
                created_at=datetime.utcnow()
            )
            mysql_session.add(job)
            mysql_session.flush()  # Flush to get the job.id
            
            # If there's a result in the SQLite data
            if row_dict.get('result'):
                import json
                result_data = row_dict['result']
                
                # If result is a JSON string, parse it
                if isinstance(result_data, str):
                    try:
                        result_json = json.loads(result_data)
                        final_analysis = result_json.get('final_analysis', result_data)
                    except:
                        final_analysis = result_data
                else:
                    final_analysis = str(result_data)
                
                result = AnalysisResult(
                    job_id=job.id,
                    final_analysis=final_analysis,
                    created_at=datetime.utcnow()
                )
                mysql_session.add(result)
    
    mysql_session.commit()
    print(f"Migrated {len(jobs)} analysis jobs")

    print("\nMigration completed successfully!")
    
    # Show summary
    print("\nDatabase summary:")
    print(f"Companies: {mysql_session.query(Company).count()}")
    print(f"Users: {mysql_session.query(User).count()}")
    print(f"Analysis Jobs: {mysql_session.query(AnalysisJob).count()}")

except Exception as e:
    print(f"Error during migration: {str(e)}")
    import traceback
    traceback.print_exc()
    mysql_session.rollback()
finally:
    sqlite_conn.close()
    mysql_session.close()