#!/usr/bin/env python3
"""
Diamond Database Sync and SKU Prefix Update Script

Copies diamond database files from Ring Legacy to SunDiamonds
and updates SKU prefixes from "RL" to "SD".

Usage:
    python sync_diamond_sku.py [--dry-run] [--no-sku-update]

Options:
    --dry-run          Show what would be done without making changes
    --no-sku-update    Copy files only, skip SKU prefix update
"""

import os
import sys
import shutil
import sqlite3
import logging
from pathlib import Path
from datetime import datetime
from typing import Optional, Tuple

# Setup logging
LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
logging.basicConfig(
    level=logging.INFO,
    format=LOG_FORMAT,
    handlers=[
        logging.StreamHandler(sys.stdout),
    ]
)
logger = logging.getLogger(__name__)

# Only update these tables (skip others even if they contain website_sku)
ALLOWED_SKU_TABLES = {
    "sku_mappings",
    "master_products",
    "fancy_sku_mappings",
    "fancy_color_diamonds",
}


class DiamondDatabaseSyncer:
    """Syncs diamond databases and updates SKU prefixes."""
    
    def __init__(
        self,
        source_dir: str = None,
        dest_dir: str = None,
        log_file: Optional[str] = None
    ):
        """
        Initialize the syncer.
        
        Args:
            source_dir: Source directory containing database files (if None, uses default)
            dest_dir: Destination directory for database files (if None, uses default)
            log_file: Optional log file path (if None, uses default)
        """
        # Get home directory to construct default paths
        home_dir = Path.home()
        
        # Set default paths relative to home directory
        if source_dir is None:
            source_dir = str(home_dir / "public_html/RINGLEGACY.COM/wp-content/uploads/diamonds")
        if dest_dir is None:
            dest_dir = str(home_dir / "public_html/sundiamonds.com/wp-content/uploads/diamonds")
        
        self.source_dir = Path(source_dir)
        self.dest_dir = Path(dest_dir)
        
        # Database files to sync
        self.db_files = [
            "lab-diamonds.db",
            "fancy-lab-diamonds.db"
        ]
        
        # Setup log file
        if log_file is None:
            self.log_file = self.dest_dir / "sync_sku_update.log"
        else:
            self.log_file = Path(log_file)
        
        # Ensure log directory exists
        try:
            self.log_file.parent.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            # If we can't create the log directory, log to a fallback location
            logger.warning(f"Could not create log directory {self.log_file.parent}: {e}")
            # Use a fallback location in the current directory or /tmp
            import tempfile
            self.log_file = Path(tempfile.gettempdir()) / "sync_sku_update.log"
            logger.warning(f"Using fallback log location: {self.log_file}")
        
        # Add file handler for logging
        if not any(isinstance(h, logging.FileHandler) for h in logger.handlers):
            try:
                file_handler = logging.FileHandler(self.log_file)
                file_handler.setFormatter(logging.Formatter(LOG_FORMAT))
                logger.addHandler(file_handler)
            except Exception as e:
                logger.warning(f"Could not create log file {self.log_file}: {e}. Logging to console only.")
    
    def ensure_destination_dir(self) -> bool:
        """Ensure destination directory exists."""
        try:
            self.dest_dir.mkdir(parents=True, exist_ok=True)
            logger.info(f"Destination directory ready: {self.dest_dir}")
            return True
        except PermissionError as e:
            logger.error(f"Permission denied creating destination directory: {e}")
            logger.error(f"Please check directory permissions: {self.dest_dir.parent}")
            return False
        except Exception as e:
            logger.error(f"Failed to create destination directory: {e}")
            return False
    
    def copy_database_file(
        self,
        filename: str,
        dry_run: bool = False,
        dest_path_override: Optional[Path] = None
    ) -> Tuple[bool, Optional[Path]]:
        """
        Copy a database file from source to destination.
        
        Args:
            filename: Name of the database file
            dry_run: If True, don't actually copy
            
        Returns:
            Tuple of (success, destination_path)
        """
        source_path = self.source_dir / filename
        dest_path = dest_path_override if dest_path_override else (self.dest_dir / filename)
        
        # Check if source exists
        if not source_path.exists():
            logger.error(f"Source file not found: {source_path}")
            return False, None
        
        if dry_run:
            source_size = source_path.stat().st_size
            logger.info(f"[DRY RUN] Would copy: {source_path} ({source_size:,} bytes) -> {dest_path}")
            return True, dest_path
        
        try:
            # Copy file
            shutil.copy2(source_path, dest_path)
            source_size = source_path.stat().st_size
            dest_size = dest_path.stat().st_size
            
            if source_size == dest_size:
                logger.info(f"✓ Copied {filename} ({source_size:,} bytes)")
                return True, dest_path
            else:
                logger.warning(
                    f"Size mismatch for {filename}: source={source_size:,} dest={dest_size:,}"
                )
                return False, dest_path
        except Exception as e:
            logger.error(f"Failed to copy {filename}: {e}")
            return False, None
    
    def get_tables_with_column(self, db_path: Path, column_name: str) -> Optional[list]:
        """
        Find all tables that contain a specific column.

        Args:
            db_path: Path to SQLite database
            column_name: Column to search for

        Returns:
            List of table names or None if not found
        """
        try:
            conn = sqlite3.connect(str(db_path))
            cursor = conn.cursor()

            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
            )
            tables = [row[0] for row in cursor.fetchall()]

            matching_tables = []
            for table_name in tables:
                if table_name not in ALLOWED_SKU_TABLES:
                    continue
                try:
                    cursor.execute(f"PRAGMA table_info({table_name})")
                    columns = [row[1] for row in cursor.fetchall()]
                    if column_name in columns:
                        matching_tables.append(table_name)
                except sqlite3.Error:
                    continue

            conn.close()

            return matching_tables if matching_tables else None
        except Exception as e:
            logger.error(f"Failed to detect tables in {db_path}: {e}")
            return None
    
    def update_sku_prefixes(
        self,
        db_path: Path,
        table_name: str,
        dry_run: bool = False,
        batch_size: int = 10000
    ) -> Tuple[bool, int]:
        """
        Update SKU prefixes from "RL" to "SD" in the database.
        
        Args:
            db_path: Path to SQLite database
            table_name: Name of the table containing SKUs
            dry_run: If True, don't actually update
            
        Returns:
            Tuple of (success, count_updated)
        """
        if not db_path.exists():
            logger.error(f"Database file not found: {db_path}")
            return False, 0
        
        try:
            conn = sqlite3.connect(str(db_path), timeout=60)
            cursor = conn.cursor()

            # Use WAL for faster bulk updates and lower lock contention
            cursor.execute("PRAGMA journal_mode=WAL")
            cursor.execute("PRAGMA synchronous=NORMAL")
            # Reduce memory footprint during large updates
            cursor.execute("PRAGMA cache_size=-20000")  # ~20MB cache
            cursor.execute("PRAGMA temp_store=FILE")
            cursor.execute("PRAGMA mmap_size=0")
            cursor.execute("PRAGMA wal_autocheckpoint=1000")
            cursor.execute("PRAGMA busy_timeout=60000")
            
            # First, check how many SKUs need updating
            cursor.execute(
                f"SELECT COUNT(*) FROM {table_name} WHERE website_sku LIKE 'RL%'"
            )
            count_to_update = cursor.fetchone()[0]
            
            if count_to_update == 0:
                logger.info(f"No SKUs with 'RL' prefix found in {db_path.name}")
                conn.close()
                return True, 0
            
            if dry_run:
                logger.info(
                    f"[DRY RUN] Would update {count_to_update} SKUs in {db_path.name}"
                )
                # Show sample SKUs that would be updated
                cursor.execute(
                    f"SELECT website_sku FROM {table_name} WHERE website_sku LIKE 'RL%' LIMIT 5"
                )
                samples = cursor.fetchall()
                if samples:
                    logger.info("Sample SKUs that would be updated:")
                    for (sku,) in samples:
                        logger.info(f"  {sku} -> {sku.replace('RL', 'SD', 1)}")
                conn.close()
                return True, count_to_update
            
            # Perform chunked updates to reduce memory pressure
            updated_total = 0
            batch_index = 0
            max_retries = 3
            retry_delay = 3
            while True:
                updated = 0
                retries = 0
                while True:
                    try:
                        cursor.execute("BEGIN IMMEDIATE")
                        cursor.execute(
                            f"""
                            UPDATE {table_name}
                            SET website_sku = 'SD' || SUBSTR(website_sku, 3)
                            WHERE rowid IN (
                                SELECT rowid FROM {table_name}
                                WHERE website_sku LIKE 'RL%'
                                LIMIT ?
                            )
                            """,
                            (batch_size,)
                        )
                        updated = cursor.rowcount or 0
                        break
                    except sqlite3.OperationalError as exc:
                        if "locked" in str(exc).lower() and retries < max_retries:
                            retries += 1
                            logger.warning(
                                f"Database locked updating {table_name}. Retrying in {retry_delay}s "
                                f"({retries}/{max_retries})..."
                            )
                            conn.rollback()
                            import time

                            time.sleep(retry_delay)
                            continue
                        raise
                if updated == 0:
                    # Ensure we don't leave an open transaction that can lock the DB
                    conn.rollback()
                    break

                updated_total += updated
                batch_index += 1
                conn.commit()

                if batch_index % 10 == 0:
                    logger.info(
                        f"Updated {updated_total:,}/{count_to_update:,} SKUs in {table_name}..."
                    )

            # Verify the update
            if conn.in_transaction:
                conn.rollback()
            cursor.execute(
                f"SELECT COUNT(*) FROM {table_name} WHERE website_sku LIKE 'RL%'"
            )
            remaining = cursor.fetchone()[0]

            if remaining == 0:
                # Flush WAL into the main database file
                cursor.execute("PRAGMA wal_checkpoint(TRUNCATE)")
                logger.info(f"✓ Updated {updated_total} SKUs in {db_path.name}")
                conn.close()
                return True, updated_total
            else:
                logger.warning(
                    f"Update incomplete: {remaining} SKUs still have 'RL' prefix in {db_path.name}"
                )
                conn.rollback()
                # Flush WAL to keep DB consistent after rollback
                cursor.execute("PRAGMA wal_checkpoint(TRUNCATE)")
                conn.close()
                return False, 0
                
        except Exception as e:
            logger.error(f"Failed to update SKUs in {db_path.name}: {e}")
            try:
                conn.rollback()
                cursor.execute("PRAGMA wal_checkpoint(TRUNCATE)")
                conn.close()
            except:
                pass
            return False, 0
    
    def finalize_database_swap(self, temp_path: Path, final_path: Path) -> bool:
        """Atomically replace the live database with the updated temp file."""
        try:
            # Clean up temp WAL/SHM files if present
            for suffix in ("-wal", "-shm"):
                sidecar = Path(str(temp_path) + suffix)
                if sidecar.exists():
                    sidecar.unlink()

            backup_path = Path(str(final_path) + ".bak")
            if final_path.exists():
                try:
                    if backup_path.exists():
                        backup_path.unlink()
                except Exception:
                    pass
                Path(final_path).replace(backup_path)

            Path(temp_path).replace(final_path)
            logger.info(f"✓ Replaced {final_path.name} (backup: {backup_path.name})")
            return True
        except Exception as e:
            logger.error(f"Failed to replace {final_path.name}: {e}")
            return False

    def sync_all(
        self,
        dry_run: bool = False,
        update_skus: bool = True
    ) -> bool:
        """
        Sync all database files and update SKU prefixes.
        
        Args:
            dry_run: If True, don't make actual changes
            update_skus: If True, update SKU prefixes after copying
            
        Returns:
            True if all operations succeeded
        """
        logger.info("=" * 60)
        logger.info("DIAMOND DATABASE SYNC AND SKU UPDATE")
        logger.info("=" * 60)
        logger.info(f"Source: {self.source_dir}")
        logger.info(f"Destination: {self.dest_dir}")
        logger.info(f"Dry Run: {dry_run}")
        logger.info(f"Update SKUs: {update_skus}")
        logger.info("=" * 60)
        
        if not self.ensure_destination_dir():
            return False
        
        success_count = 0
        total_count = len(self.db_files)
        
        for filename in self.db_files:
            logger.info(f"\nProcessing: {filename}")

            dest_path = self.dest_dir / filename
            temp_path = self.dest_dir / f"{filename}.updating"

            # Copy the file (use temp path when updating SKUs)
            copy_target = temp_path if (update_skus and not dry_run) else dest_path
            success, copied_path = self.copy_database_file(
                filename,
                dry_run=dry_run,
                dest_path_override=copy_target
            )
            if not success or copied_path is None:
                logger.error(f"Failed to copy {filename}")
                continue
            
            # Update SKU prefixes if requested
            if update_skus and not dry_run:
                # Detect tables with website_sku column
                table_names = self.get_tables_with_column(copy_target, "website_sku")
                if table_names:
                    logger.info(f"Detected tables with website_sku: {', '.join(table_names)}")
                    table_success = True
                    total_updated = 0
                    for table_name in table_names:
                        sku_success, count = self.update_sku_prefixes(
                            copy_target, table_name, dry_run=False
                        )
                        total_updated += count
                        if not sku_success:
                            table_success = False
                    if table_success:
                        logger.info(f"✓ Updated {total_updated} SKUs across {len(table_names)} tables")
                        if copy_target != dest_path:
                            if not self.finalize_database_swap(copy_target, dest_path):
                                logger.error(f"Failed to replace {dest_path.name} after update")
                                continue
                        success_count += 1
                    else:
                        logger.error(f"SKU update failed for {filename}")
                else:
                    logger.warning(
                        f"Could not detect tables with website_sku in {filename}, skipping SKU update"
                    )
                    success_count += 1  # Still count as success if copy worked
            elif update_skus and dry_run:
                # In dry run, check the source file
                source_path = self.source_dir / filename
                if source_path.exists():
                    table_names = self.get_tables_with_column(source_path, "website_sku")
                    if table_names:
                        for table_name in table_names:
                            self.update_sku_prefixes(source_path, table_name, dry_run=True)
                        success_count += 1  # Count as success in dry run
                else:
                    logger.warning(f"Source file not found for SKU check: {source_path}")
            else:
                # Copy only, no SKU update
                success_count += 1
        
        logger.info("\n" + "=" * 60)
        if success_count == total_count:
            logger.info(f"✓ SUCCESS: {success_count}/{total_count} databases synced")
            logger.info("=" * 60)
            return True
        else:
            logger.error(f"✗ PARTIAL: {success_count}/{total_count} databases synced")
            logger.info("=" * 60)
            return False


def main():
    """Main entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Sync diamond databases and update SKU prefixes from RL to SD"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be done without making changes"
    )
    parser.add_argument(
        "--no-sku-update",
        action="store_true",
        help="Copy files only, skip SKU prefix update"
    )
    parser.add_argument(
        "--source-dir",
        default=None,
        help="Source directory (default: ~/public_html/RINGLEGACY.COM/wp-content/uploads/diamonds)"
    )
    parser.add_argument(
        "--dest-dir",
        default=None,
        help="Destination directory (default: ~/public_html/sundiamonds.com/wp-content/uploads)"
    )
    parser.add_argument(
        "--log-file",
        help="Custom log file path (default: {dest_dir}/sync_sku_update.log)"
    )
    
    args = parser.parse_args()
    
    try:
        syncer = DiamondDatabaseSyncer(
            source_dir=args.source_dir,
            dest_dir=args.dest_dir,
            log_file=args.log_file
        )
        
        success = syncer.sync_all(
            dry_run=args.dry_run,
            update_skus=not args.no_sku_update
        )
        
        sys.exit(0 if success else 1)
        
    except KeyboardInterrupt:
        logger.info("\nOperation cancelled by user")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()