#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import logging
import sys
from typing import Dict, Tuple, Any
from db.functions import settype, dbmongo
from db.schema_generator import MongoToPostgresSchemaGenerator, MongoDocumentInserter
import settings.config


# Configuration constants
SAMPLE_SIZE = 10000
BATCH_SIZE = 10000
SCHEMA_NAME = 'updates'
PK_FIELD = 'id'

# Table-to-collection mapping
TABLE_CONFIG = {
    'seriesdata': 'mgseries',
    'episodesdata': 'mgepisodes',
    'actorsdata': 'mgactors',
    'charactersdata': 'mgcharacters',
    'crewdata': 'mgcrew',
}


def setup_logger() -> logging.Logger:
    """Initialize and return configured logger."""
    logger = logging.getLogger(__name__)
    if not logger.handlers:
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.setLevel(logging.INFO)
    return logger


def initialize_environment(root_dir: str) -> Tuple[Dict[str, Any], Any, Any]:
    """Initialize config, database engine, and MongoDB updater.
    
    Args:
        root_dir: Root directory path
        
    Returns:
        Tuple of (options dict, db engine, mongo updater)
    """
    logger = logging.getLogger(__name__)
    try:
        config_options = settings.config.Config(root_dir)
        options = config_options.config_options
        
        if not options:
            raise ValueError("Configuration is empty")
            
        dbtype = options.get('dbtype')
        apitype = options.get('apitype')
        
        if not dbtype or not apitype:
            raise ValueError("Missing required config options: dbtype, apitype")
        
        dbs = settype(dbtype, apitype, options)
        dbengine = dbs.dbengine
        mongo_updater = dbmongo(options)
        
        logger.info("Environment initialized successfully")
        return options, dbengine, mongo_updater
        
    except Exception as e:
        logger.error(f"Failed to initialize environment: {e}")
        raise


def create_tables(
    generator: MongoToPostgresSchemaGenerator,
    mongo_updater: Any,
    dbengine: Dict[str, Any],
    table_config: Dict[str, str],
    logger: logging.Logger
) -> Dict[str, Dict[str, Any]]:
    """Create PostgreSQL tables and return cached schemas.
    
    Args:
        generator: Schema generator instance
        mongo_updater: MongoDB updater instance
        dbengine: Database engine configuration
        table_config: Mapping of table names to MongoDB collection attributes
        logger: Logger instance
        
    Returns:
        Dict mapping table names to their schemas
    """
    schemas = {}
    
    for table_name, collection_attr in table_config.items():
        try:
            # Get MongoDB collection
            collection = getattr(mongo_updater, collection_attr, None)
            if not collection:
                logger.warning(f"Collection '{collection_attr}' not found on mongo_updater")
                continue
            
            # Analyze schema
            schema = generator.analyze_collection(mongo_updater, collection)
            schemas[table_name] = schema
            
            # Generate and display SQL
            sql = generator.generate_create_table_sql(
                table_name,
                schema_name=SCHEMA_NAME,
                pk_field=PK_FIELD,
            )
            logger.info(f"Schema for {table_name}:\n{sql}")
            
            # Create table in PostgreSQL
            generator.create_table_in_postgres(
                dbengine,
                table_name,
                schema_name=SCHEMA_NAME,
                pk_field=PK_FIELD,
                drop_existing=True,
            )
            logger.info(f"Table '{table_name}' created/recreated successfully")
            
        except Exception as e:
            logger.error(f"Failed to create table '{table_name}': {e}")
            continue
    
    return schemas


def build_upsert_clause(schema: Dict[str, Any], pk_field: str) -> str:
    """Build SQL ON CONFLICT clause for upsert.
    
    Args:
        schema: Column schema dictionary
        pk_field: Primary key field name
        
    Returns:
        ON CONFLICT clause string
    """
    # Exclude primary key and MongoDB's _id field from updates
    update_cols = [
        col.lower() for col in schema.keys()
        if col.lower() != pk_field.lower() and col != '_id'
    ]
    
    if not update_cols:
        return f"ON CONFLICT ({pk_field}) DO NOTHING"
    
    update_clause = ', '.join([f"{col}=EXCLUDED.{col}" for col in update_cols])
    return f"ON CONFLICT ({pk_field}) DO UPDATE SET {update_clause}"


def should_update_document(
    doc_id: Any,
    doc_updated: Any,
    dbengine: Dict[str, Any],
    table_name: str,
    schema_name: str,
    logger: logging.Logger
) -> bool:
    """Compare document's updated timestamp to the database row's updated timestamp.
    
    Args:
        doc_id: Document ID from MongoDB
        doc_updated: Updated timestamp from MongoDB document
        dbengine: Database engine configuration
        table_name: Target table name
        schema_name: Target schema name
        logger: Logger instance
        
    Returns:
        True if document should be updated (MongoDB is newer), False otherwise
    """
    if doc_updated is None:
        logger.debug(f"Document {doc_id} has no 'updated' field, will update")
        return True
    
    try:
        from sqlalchemy import text
        
        # Query the database for the existing record's updated timestamp
        query = text(f"""
            SELECT updated FROM {schema_name}.{table_name}
            WHERE id = :id
        """)
        
        with dbengine['engine'].connect() as conn:
            result = conn.execute(query, {'id': doc_id}).fetchone()
        
        # If record doesn't exist, update it
        if result is None:
            logger.debug(f"Record {doc_id} not in DB, will insert")
            return True
        
        db_updated = result[0]
        
        # Compare timestamps
        if db_updated is None:
            logger.debug(f"Record {doc_id} has NULL updated in DB, will update")
            return True
        
        # Only update if MongoDB version is newer
        should_update = doc_updated > db_updated
        
        if not should_update:
            logger.debug(f"Record {doc_id}: DB updated={db_updated} >= doc updated={doc_updated}, skipping")
        else:
            logger.debug(f"Record {doc_id}: doc updated={doc_updated} > DB updated={db_updated}, updating")
        
        return should_update
        
    except Exception as e:
        logger.warning(f"Error comparing timestamps for {doc_id}: {e}. Proceeding with update.")
        return True


def insert_documents(
    inserter: MongoDocumentInserter,
    mongo_updater: Any,
    dbengine: Dict[str, Any],
    table_config: Dict[str, str],
    schemas: Dict[str, Dict[str, Any]],
    logger: logging.Logger,
    compare_timestamps: bool = True
) -> None:
    """Insert/upsert documents from MongoDB collections to PostgreSQL.
    
    Args:
        inserter: Document inserter instance
        mongo_updater: MongoDB updater instance
        dbengine: Database engine configuration
        table_config: Mapping of table names to MongoDB collection attributes
        schemas: Cached schemas from table creation
        logger: Logger instance
        compare_timestamps: If True, only update if MongoDB doc is newer than DB record
    """
    total_inserted = 0
    total_skipped = 0
    
    for table_name, collection_attr in table_config.items():
        if table_name not in schemas:
            logger.warning(f"Skipping '{table_name}' - schema not available")
            continue
        
        try:
            # Get collection and schema
            collection = getattr(mongo_updater, collection_attr)
            schema = schemas[table_name]
            
            # Build upsert clause
            on_conflict = build_upsert_clause(schema, PK_FIELD)
            
            # If timestamp comparison is enabled, filter documents before inserting
            if compare_timestamps:
                logger.info(f"Filtering {table_name} documents by updated timestamp and batching every {BATCH_SIZE}...")
                batch = []
                batch_num = 0

                for doc in collection.find():
                    doc_id = doc.get(PK_FIELD)
                    doc_updated = doc.get('updated')

                    if should_update_document(doc_id, doc_updated, dbengine, table_name, SCHEMA_NAME, logger):
                        batch.append(doc)
                    else:
                        total_skipped += 1

                    # When batch reaches size, insert and reset
                    if len(batch) >= BATCH_SIZE:
                        batch_num += 1
                        try:
                            # Explicit transaction per batch
                            with dbengine['engine'].begin():
                                count = inserter.insert_documents(
                                    engine=dbengine,
                                    table_name=table_name,
                                    documents=batch,
                                    schema_name=SCHEMA_NAME,
                                    on_conflict=on_conflict,
                                )
                            total_inserted += count
                            logger.info(f"Batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)")
                        except Exception as e:
                            logger.error(f"Batch {batch_num} failed for '{table_name}': {e}")
                        finally:
                            batch = []

                # Insert any remaining docs in the final batch
                if batch:
                    batch_num += 1
                    try:
                        with dbengine['engine'].begin():
                            count = inserter.insert_documents(
                                engine=dbengine,
                                table_name=table_name,
                                documents=batch,
                                schema_name=SCHEMA_NAME,
                                on_conflict=on_conflict,
                            )
                        total_inserted += count
                        logger.info(f"Final batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)")
                    except Exception as e:
                        logger.error(f"Final batch {batch_num} failed for '{table_name}': {e}")
                    finally:
                        batch = []
            else:
                # Insert all documents without timestamp comparison
                # Stream and batch inserts to avoid memory spikes
                batch = []
                batch_num = 0
                for doc in collection.find():
                    batch.append(doc)
                    if len(batch) >= BATCH_SIZE:
                        batch_num += 1
                        try:
                            with dbengine['engine'].begin():
                                count = inserter.insert_documents(
                                    engine=dbengine,
                                    table_name=table_name,
                                    documents=batch,
                                    schema_name=SCHEMA_NAME,
                                    on_conflict=on_conflict,
                                )
                            total_inserted += count
                            logger.info(f"Batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)")
                        except Exception as e:
                            logger.error(f"Batch {batch_num} failed for '{table_name}': {e}")
                        finally:
                            batch = []

                if batch:
                    batch_num += 1
                    try:
                        with dbengine['engine'].begin():
                            count = inserter.insert_documents(
                                engine=dbengine,
                                table_name=table_name,
                                documents=batch,
                                schema_name=SCHEMA_NAME,
                                on_conflict=on_conflict,
                            )
                        total_inserted += count
                        logger.info(f"Final batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)")
                    except Exception as e:
                        logger.error(f"Final batch {batch_num} failed for '{table_name}': {e}")
                    finally:
                        batch = []
            
        except Exception as e:
            logger.error(f"Failed to insert documents into '{table_name}': {e}")
            continue
    
    logger.info(f"Total documents inserted/updated: {total_inserted}, skipped: {total_skipped}")


def main() -> int:
    """Main entry point for MongoDB to PostgreSQL sync."""
    logger = setup_logger()
    
    try:
        root_dir = os.getcwd()
        
        # Initialize environment
        options, dbengine, mongo_updater = initialize_environment(root_dir)
        
        # Create schema generator and inserter
        generator = MongoToPostgresSchemaGenerator(sample_size=SAMPLE_SIZE)
        inserter = MongoDocumentInserter(batch_size=BATCH_SIZE)
        
        # Create tables and get cached schemas
        logger.info("Creating PostgreSQL tables...")
        schemas = create_tables(generator, mongo_updater, dbengine, TABLE_CONFIG, logger)
        
        if not schemas:
            logger.error("No schemas were successfully created. Aborting insert.")
            return 1
        
        # Insert documents from MongoDB
        logger.info("Inserting documents from MongoDB...")
        insert_documents(inserter, mongo_updater, dbengine, TABLE_CONFIG, schemas, logger)
        
        logger.info("MongoDB to PostgreSQL sync completed successfully")
        return 0
        
    except Exception as e:
        logger.error(f"Sync failed: {e}", exc_info=True)
        return 1


if __name__ == "__main__":
    sys.exit(main())
