"""
SQLite to MySQL Migration Script

This script migrates data from local SQLite databases (`trading_history.db` and `data/market_data.db`)
to a remote MySQL database defined in environment variables.

Environment Variables required in .env:
- MYSQL_HOST
- MYSQL_PORT (default: 3306)
- MYSQL_USER
- MYSQL_PASSWORD
- MYSQL_DB

Usage:
    python scripts/migrate_sqlite_to_mysql.py
"""

import sys
import os
import sqlite3
import pymysql
from dotenv import load_dotenv
from tqdm import tqdm
import pandas as pd

# Add project root to sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

load_dotenv()

def get_mysql_connection():
    return pymysql.connect(
        host=os.getenv('MYSQL_HOST'),
        port=int(os.getenv('MYSQL_PORT', 3306)),
        user=os.getenv('MYSQL_USER'),
        password=os.getenv('MYSQL_PASSWORD'),
        database=os.getenv('MYSQL_DB'),
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

def create_mysql_tables(conn):
    with conn.cursor() as cursor:
        # Create trades table
        print("Creating 'trades' table in MySQL...")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS trades (
                id INT AUTO_INCREMENT PRIMARY KEY,
                timestamp DATETIME,
                market_title VARCHAR(255),
                market_slug VARCHAR(255),
                model_name VARCHAR(50),
                prediction VARCHAR(10),
                probability DOUBLE,
                action VARCHAR(10),
                price DOUBLE,
                amount DOUBLE,
                size DOUBLE,
                status VARCHAR(20),
                error_message TEXT,
                raw_data JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                INDEX idx_market_slug (market_slug),
                INDEX idx_status (status)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        """)

        # Create ohlcv table
        print("Creating 'ohlcv' table in MySQL...")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS ohlcv (
                timestamp DATETIME,
                symbol VARCHAR(20),
                timeframe VARCHAR(10),
                open DOUBLE,
                high DOUBLE,
                low DOUBLE,
                close DOUBLE,
                volume DOUBLE,
                PRIMARY KEY (symbol, timeframe, timestamp),
                INDEX idx_time (timestamp)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        """)
    conn.commit()

def migrate_trades(mysql_conn):
    sqlite_db_path = "trading_history.db"
    if not os.path.exists(sqlite_db_path):
        print(f"Skipping trades migration: {sqlite_db_path} not found.")
        return

    print(f"Migrating trades from {sqlite_db_path}...")
    
    with sqlite3.connect(sqlite_db_path) as sqlite_conn:
        sqlite_cursor = sqlite_conn.cursor()
        
        # Check if table exists
        sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='trades';")
        if not sqlite_cursor.fetchone():
            print("Table 'trades' not found in SQLite.")
            return

        sqlite_cursor.execute("SELECT * FROM trades")
        rows = sqlite_cursor.fetchall()
        columns = [description[0] for description in sqlite_cursor.description]
        
        if not rows:
            print("No trades data to migrate.")
            return

        print(f"Found {len(rows)} trade records.")
        
        # Insert into MySQL
        with mysql_conn.cursor() as mysql_cursor:
            # Prepare SQL statement
            # Map columns. SQLite 'timestamp' is text, MySQL needs valid datetime format or string.
            # We skip 'id' to let MySQL handle auto-increment or we can preserve it.
            # Let's preserve ID if possible, but usually better to let destination handle it unless we need consistency.
            # User wants migration, preserving ID is safer for history.
            
            placeholders = ', '.join(['%s'] * len(columns))
            cols_str = ', '.join(columns)
            
            sql = f"INSERT IGNORE INTO trades ({cols_str}) VALUES ({placeholders})"
            
            mysql_cursor.executemany(sql, rows)
            mysql_conn.commit()
            print(f"Migrated {mysql_cursor.rowcount} trade records.")

def migrate_ohlcv(mysql_conn):
    sqlite_db_path = "data/market_data.db"
    if not os.path.exists(sqlite_db_path):
        print(f"Skipping OHLCV migration: {sqlite_db_path} not found.")
        return

    print(f"Migrating OHLCV from {sqlite_db_path}...")
    
    with sqlite3.connect(sqlite_db_path) as sqlite_conn:
        # Use pandas for easier chunk reading
        try:
            # Check total count
            total_rows = pd.read_sql_query("SELECT COUNT(*) FROM ohlcv", sqlite_conn).iloc[0, 0]
            print(f"Found {total_rows} OHLCV records.")
            
            chunk_size = 10000
            offset = 0
            
            with tqdm(total=total_rows, desc="Migrating OHLCV") as pbar:
                while True:
                    df = pd.read_sql_query(f"SELECT * FROM ohlcv LIMIT {chunk_size} OFFSET {offset}", sqlite_conn)
                    if df.empty:
                        break
                    
                    # Convert DataFrame to list of tuples
                    data = [tuple(x) for x in df.to_numpy()]
                    
                    with mysql_conn.cursor() as mysql_cursor:
                        sql = """
                            INSERT IGNORE INTO ohlcv 
                            (timestamp, symbol, timeframe, open, high, low, close, volume) 
                            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                        """
                        # Select columns explicitly to match SQL
                        cols = ['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']
                        batch_data = [tuple(row[col] for col in cols) for _, row in df.iterrows()]
                        
                        mysql_cursor.executemany(sql, batch_data)
                        mysql_conn.commit()
                    
                    offset += chunk_size
                    pbar.update(len(df))
                    
        except Exception as e:
            print(f"Error reading/migrating OHLCV: {e}")

def main():
    if not os.getenv('MYSQL_HOST'):
        print("Error: MYSQL_HOST not set in .env")
        return

    try:
        print("Connecting to MySQL...")
        mysql_conn = get_mysql_connection()
        print("Connected.")
        
        create_mysql_tables(mysql_conn)
        
        migrate_trades(mysql_conn)
        migrate_ohlcv(mysql_conn)
        
        mysql_conn.close()
        print("\nMigration completed successfully.")
        
    except Exception as e:
        print(f"\nMigration failed: {e}")

if __name__ == "__main__":
    main()