import subprocess
import sys
import os
import re
import argparse
import gzip
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv, find_dotenv

# Load environment variables from .env file
if not find_dotenv():
    print("Error: .env file not found.")
    sys.exit(1)
load_dotenv()

MYSQL_HOST = os.getenv('MYSQL_HOST')
MYSQL_USER = os.getenv('MYSQL_USER')
MYSQL_PASSWORD = os.getenv('MYSQL_PASSWORD')
BACKUP_DIR = os.getenv('BACKUP_DIR')

# Check if any required environment variable is missing
missing = [name for name, var in [
    ('MYSQL_HOST', MYSQL_HOST),
    ('MYSQL_USER', MYSQL_USER),
    ('MYSQL_PASSWORD', MYSQL_PASSWORD),
    ('BACKUP_DIR', BACKUP_DIR)
] if var is None]
if missing:
    print(f"Error: Missing environment variables: {', '.join(missing)}")
    sys.exit(1)

# Type assertions for type checkers
assert MYSQL_HOST is not None
assert MYSQL_USER is not None
assert MYSQL_PASSWORD is not None
assert BACKUP_DIR is not None


def find_backups(database, backup_dir):
    """Find all backup files for the specified database."""
    if not os.path.isdir(backup_dir):
        print(f"Error: Backup directory does not exist: {backup_dir}")
        return []

    # Sanitize database name same way backup script does
    safe_db_name = database.replace('/', '_').replace('\\', '_').replace(' ', '_')

    # Match exact db name followed by timestamp: {name}_YYYY-MM-DD_HHMMSS.sql[.gz]
    # Prevents e.g. "scike" matching "scike_guitar_*" backups.
    pattern = re.compile(rf"^{re.escape(safe_db_name)}_\d{{4}}-\d{{2}}-\d{{2}}_\d{{6}}\.sql(\.gz)?$")
    backups = []
    for file in os.listdir(backup_dir):
        if pattern.match(file):
            file_path = os.path.join(backup_dir, file)
            # Extract timestamp from filename (format: dbname_YYYY-MM-DD_HHMMSS.sql[.gz])
            try:
                # Get file stats
                stat = os.stat(file_path)
                backups.append({
                    'path': file_path,
                    'filename': file,
                    'size': stat.st_size,
                    'mtime': stat.st_mtime
                })
            except Exception as e:
                print(f"Warning: Could not read file {file}: {e}")
                continue

    # Sort by modification time (newest first)
    backups.sort(key=lambda x: x['mtime'], reverse=True)
    return backups


def format_size(size_bytes):
    """Format file size in human-readable format."""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.2f} KB"
    else:
        return f"{size_bytes / (1024 * 1024):.2f} MB"


def format_age(mtime):
    """Format file age in human-readable format."""
    age_seconds = datetime.now().timestamp() - mtime
    age_hours = age_seconds / 3600

    if age_hours < 1:
        return f"{int(age_seconds / 60)} minutes ago"
    elif age_hours < 24:
        return f"{int(age_hours)} hours ago"
    else:
        age_days = int(age_hours / 24)
        return f"{age_days} day{'s' if age_days > 1 else ''} ago"


def display_backup_options(backups, limit=5):
    """Display backup options and return user selection."""
    print("\n" + "=" * 70)
    print(f"Found {len(backups)} backup(s)")
    print("=" * 70)

    display_count = min(limit, len(backups))
    for i, backup in enumerate(backups[:display_count], 1):
        timestamp = datetime.fromtimestamp(backup['mtime']).strftime('%Y-%m-%d %H:%M:%S')
        size = format_size(backup['size'])
        age = format_age(backup['mtime'])
        compressed = " [compressed]" if backup['filename'].endswith('.gz') else ""
        print(f"{i}. {backup['filename']}")
        print(f"   Date: {timestamp} ({age})")
        print(f"   Size: {size}{compressed}")
        print()

    if len(backups) > display_count:
        print(f"... and {len(backups) - display_count} more backup(s)\n")

    print("=" * 70)
    print(f"Press Enter to use most recent backup, or enter 1-{display_count}")
    if len(backups) > display_count:
        print(f"Enter 'all' to see all backups, or 'custom' to specify a file path")
    else:
        print("Enter 'custom' to specify a custom file path")
    print("=" * 70)

    choice = input("\nYour choice [Enter for most recent]: ").strip().lower()

    if choice == '':
        return backups[0]['path']
    elif choice == 'all':
        return display_backup_options(backups, limit=len(backups))
    elif choice == 'custom':
        custom_path = input("Enter full path to backup file: ").strip()
        if not os.path.isfile(custom_path):
            print(f"Error: File not found: {custom_path}")
            sys.exit(1)
        return custom_path
    else:
        try:
            index = int(choice) - 1
            if 0 <= index < display_count:
                return backups[index]['path']
            else:
                print(f"Error: Invalid choice. Please enter 1-{display_count}")
                sys.exit(1)
        except ValueError:
            print("Error: Invalid input. Please enter a number or 'custom'")
            sys.exit(1)


def confirm_restore(db_name, file_path):
    """Ask user to confirm before dropping and restoring database."""
    file_size = format_size(os.path.getsize(file_path))
    file_time = datetime.fromtimestamp(os.path.getmtime(file_path)).strftime('%Y-%m-%d %H:%M:%S')

    print("\n" + "!" * 70)
    print("WARNING: This will DROP the existing database and restore from backup")
    print("!" * 70)
    print(f"Database: {db_name}")
    print(f"Backup file: {os.path.basename(file_path)}")
    print(f"Backup date: {file_time}")
    print(f"File size: {file_size}")
    print("!" * 70)

    response = input("\nType 'yes' to continue or anything else to cancel: ").strip().lower()
    if response != 'yes':
        print("Restore cancelled by user.")
        sys.exit(0)


def drop_database(db_name):
    """Drop the database if it exists."""
    env = os.environ.copy()
    env['MYSQL_PWD'] = MYSQL_PASSWORD

    try:
        drop_db_command = [
            'mysql',
            f'--host={MYSQL_HOST}',
            f'--user={MYSQL_USER}',
            '-e',
            f'DROP DATABASE IF EXISTS {db_name};'
        ]
        subprocess.run(drop_db_command, check=True, env=env,
                      stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
        print(f"✓ Database '{db_name}' dropped successfully")
    except subprocess.CalledProcessError as e:
        print(f"Error dropping database: {e}")
        if e.stderr:
            print(f"MySQL error: {e.stderr}")
        sys.exit(1)


def create_database(db_name):
    """Create a new empty database."""
    env = os.environ.copy()
    env['MYSQL_PWD'] = MYSQL_PASSWORD

    try:
        create_db_command = [
            'mysql',
            f'--host={MYSQL_HOST}',
            f'--user={MYSQL_USER}',
            '-e',
            f'CREATE DATABASE {db_name};'
        ]
        subprocess.run(create_db_command, check=True, env=env,
                      stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
        print(f"✓ Database '{db_name}' created successfully")
    except subprocess.CalledProcessError as e:
        print(f"Error creating database: {e}")
        if e.stderr:
            print(f"MySQL error: {e.stderr}")
        sys.exit(1)


def load_database_from_file(db_name, file_path):
    """Load database from SQL file (handles both .sql and .sql.gz files)."""
    env = os.environ.copy()
    env['MYSQL_PWD'] = MYSQL_PASSWORD

    is_compressed = file_path.endswith('.gz')

    try:
        print(f"✓ Loading database from backup {('(decompressing)' if is_compressed else '')}...")
        start_time = datetime.now()

        if is_compressed:
            # Use gunzip to decompress and pipe to mysql
            # This is more reliable than Python's gzip module with subprocess
            gunzip_command = ['gunzip', '-c', file_path]
            mysql_command = [
                'mysql',
                '--binary-mode',
                f'--host={MYSQL_HOST}',
                f'--user={MYSQL_USER}',
                db_name
            ]

            # Start gunzip process
            gunzip_process = subprocess.Popen(gunzip_command, stdout=subprocess.PIPE,
                                              stderr=subprocess.PIPE, env=env)

            # Pipe gunzip output to mysql
            mysql_process = subprocess.Popen(mysql_command, stdin=gunzip_process.stdout,
                                            stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)

            # Close gunzip stdout in parent to allow gunzip to receive SIGPIPE if mysql exits
            gunzip_process.stdout.close()

            # Wait for mysql to complete and get output
            mysql_stdout, mysql_stderr = mysql_process.communicate()

            # Check for errors
            if mysql_process.returncode != 0:
                error_msg = mysql_stderr.decode('utf-8', errors='replace') if mysql_stderr else 'Unknown error'
                raise subprocess.CalledProcessError(mysql_process.returncode, mysql_command,
                                                   output=mysql_stdout, stderr=mysql_stderr)

        else:
            # Read regular SQL file and pipe to mysql
            mysql_command = [
                'mysql',
                '--binary-mode',
                f'--host={MYSQL_HOST}',
                f'--user={MYSQL_USER}',
                db_name
            ]

            with open(file_path, 'rb') as f:
                result = subprocess.run(mysql_command, stdin=f, check=True, env=env,
                                      stderr=subprocess.PIPE, stdout=subprocess.PIPE)

        elapsed = (datetime.now() - start_time).total_seconds()
        print(f"✓ Database loaded successfully in {elapsed:.2f} seconds")

    except subprocess.CalledProcessError as e:
        print(f"Error loading database from file: {e}")
        if e.stderr:
            # Decode stderr for display (mysql error messages are usually UTF-8)
            error_msg = e.stderr.decode('utf-8', errors='replace') if isinstance(e.stderr, bytes) else e.stderr
            print(f"MySQL error: {error_msg}")
        sys.exit(1)
    except Exception as e:
        print(f"Unexpected error loading database: {e}")
        sys.exit(1)


def main():
    # Parse arguments
    parser = argparse.ArgumentParser(
        description='Restore a MySQL database from backup file.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Interactive: choose from recent backups
  python db_restore.py mydatabase

  # Use specific backup file
  python db_restore.py mydatabase --file /path/to/backup.sql.gz

  # Non-interactive: use most recent backup without prompting
  python db_restore.py mydatabase --yes

  # Restore from a different backup directory
  python db_restore.py mydatabase --backup-dir /path/to/backups --yes
        """)
    parser.add_argument('database', type=str,
                        help='Name of the database to restore')
    parser.add_argument('--file', type=str, dest='sqlfile',
                        help='Path to specific SQL file (optional, will auto-discover if not provided)')
    parser.add_argument('--backup-dir', type=str,
                        help='Override backup directory for finding backups (default: BACKUP_DIR from .env)')
    parser.add_argument('--yes', '-y', action='store_true',
                        help='Auto-confirm restore (use most recent backup without prompting)')
    args = parser.parse_args()

    db_name = args.database

    # Use command-line backup dir if provided, otherwise use env var
    backup_dir = args.backup_dir if args.backup_dir else BACKUP_DIR

    # Determine which backup file to use
    if args.sqlfile:
        # User specified a file
        sql_file_path = args.sqlfile
        if not os.path.isfile(sql_file_path):
            print(f"Error: File not found: {sql_file_path}")
            sys.exit(1)
    else:
        # Auto-discover backups
        print(f"Searching for backups of '{db_name}' in {backup_dir}...")
        backups = find_backups(db_name, backup_dir)

        if not backups:
            print(f"Error: No backups found for database '{db_name}' in {backup_dir}")
            sys.exit(1)

        if args.yes:
            # Non-interactive mode: use most recent
            sql_file_path = backups[0]['path']
            print(f"Using most recent backup: {os.path.basename(sql_file_path)}")
        else:
            # Interactive mode: let user choose
            sql_file_path = display_backup_options(backups)

    # Confirm before proceeding (unless --yes flag)
    if not args.yes:
        confirm_restore(db_name, sql_file_path)

    # Perform restore
    print("\nStarting restore process...")
    print("-" * 70)
    drop_database(db_name)
    create_database(db_name)
    load_database_from_file(db_name, sql_file_path)

    # Success summary
    print("-" * 70)
    print("\n" + "=" * 70)
    print("RESTORE COMPLETED SUCCESSFULLY")
    print("=" * 70)
    print(f"Database: {db_name}")
    print(f"Restored from: {os.path.basename(sql_file_path)}")
    print("=" * 70)


if __name__ == "__main__":
    main()