import asyncio
import aiosqlite
import aiomysql
from datetime import datetime
from zoneinfo import ZoneInfo
import os
from dotenv import load_dotenv

load_dotenv()

# MySQL connection settings
MYSQL_USER = os.getenv("MYSQL_USER", "bot")
MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "secret")
MYSQL_HOST = os.getenv("MYSQL_HOST", "localhost")
MYSQL_PORT = os.getenv("MYSQL_PORT", "3306")
MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "vpn_shop")

async def check_tables_exist(cursor):
    """Check if all required tables exist"""
    tables = ['users', 'vpn_plans', 'vpn_subscriptions', 'payments']
    for table in tables:
        await cursor.execute(f"SHOW TABLES LIKE '{table}'")
        if not await cursor.fetchone():
            raise Exception(f"Table '{table}' does not exist. Please run init_mysql.py first.")

async def migrate_data():
    print("Starting migration process...")
    
    # Connect to MySQL first to check tables
    mysql_conn = await aiomysql.connect(
        host=MYSQL_HOST,
        port=int(MYSQL_PORT),
        user=MYSQL_USER,
        password=MYSQL_PASSWORD,
        db=MYSQL_DATABASE,
        autocommit=True
    )
    mysql_cur = await mysql_conn.cursor()

    try:
        # Check if tables exist
        print("Checking if tables exist...")
        await check_tables_exist(mysql_cur)
        print("All required tables found!")

        # Connect to SQLite database
        print("Connecting to SQLite database...")
        async with aiosqlite.connect("vpn_shop.db") as sqlite_conn:
            sqlite_conn.row_factory = aiosqlite.Row
            
            try:
                # Migrate users
                print("Migrating users...")
                async with sqlite_conn.execute("SELECT * FROM users") as cursor:
                    users = await cursor.fetchall()
                    for user in users:
                        await mysql_cur.execute(
                            "INSERT INTO users (telegram_id, username, created_at, ban, ref_id) "
                            "VALUES (%s, %s, %s, %s, %s)",
                            (
                                user['telegram_id'],
                                user['username'],
                                user['created_at'],
                                user['ban'],
                                user['ref_id']
                            )
                        )
                print(f"Migrated {len(users)} users")
                
                # Migrate VPN plans
                print("Migrating VPN plans...")
                async with sqlite_conn.execute("SELECT * FROM vpn_plans") as cursor:
                    plans = await cursor.fetchall()
                    for plan in plans:
                        await mysql_cur.execute(
                            "INSERT INTO vpn_plans (id, name, duration_days, price, description) "
                            "VALUES (%s, %s, %s, %s, %s)",
                            (
                                plan['id'],
                                plan['name'],
                                plan['duration_days'],
                                plan['price'],
                                plan['description']
                            )
                        )
                print(f"Migrated {len(plans)} VPN plans")
                
                # Migrate VPN subscriptions
                print("Migrating VPN subscriptions...")
                async with sqlite_conn.execute("SELECT * FROM vpn_subscriptions") as cursor:
                    subscriptions = await cursor.fetchall()
                    for sub in subscriptions:
                        await mysql_cur.execute(
                            "INSERT INTO vpn_subscriptions "
                            "(id, user_id, config_data, start_date, end_date, plan_id, server_name, status) "
                            "VALUES (%s, %s, %s, %s, %s, %s, %s, %s)",
                            (
                                sub['id'],
                                sub['user_id'],
                                sub['config_data'],
                                sub['start_date'],
                                sub['end_date'],
                                sub['plan_id'],
                                sub['server_name'],
                                sub['status']
                            )
                        )
                print(f"Migrated {len(subscriptions)} VPN subscriptions")
                
                # Migrate payments
                print("Migrating payments...")
                async with sqlite_conn.execute("SELECT * FROM payments") as cursor:
                    payments = await cursor.fetchall()
                    for payment in payments:
                        await mysql_cur.execute(
                            "INSERT INTO payments "
                            "(id, user_id, amount, payment_id, status, created_at, paid_at) "
                            "VALUES (%s, %s, %s, %s, %s, %s, %s)",
                            (
                                payment['id'],
                                payment['user_id'],
                                payment['amount'],
                                payment['payment_id'],
                                payment['status'],
                                payment['created_at'],
                                payment['paid_at']
                            )
                        )
                print(f"Migrated {len(payments)} payments")
                
                print("Migration completed successfully!")
                
            except Exception as e:
                print(f"Error during migration: {e}")
                raise
                
    except Exception as e:
        print(f"Error: {e}")
        raise
    finally:
        await mysql_cur.close()
        mysql_conn.close()

if __name__ == "__main__":
    asyncio.run(migrate_data()) 