"""Generate PostgreSQL table schemas from MongoDB document structures."""

__author__ = 'Wendell Jones'

from typing import Any, Dict, List, Optional, Set
from pymongo import MongoClient
from sqlalchemy import text, MetaData, insert, Table
import logging
import json

logger = logging.getLogger(__name__)


class MongoToPostgresSchemaGenerator:
    """Infer PostgreSQL schema from MongoDB collection and create tables."""

    # Map MongoDB/Python types to PostgreSQL types
    TYPE_MAPPING = {
        'null': 'TEXT',
        'boolean': 'BOOLEAN',
        'integer': 'BIGINT',
        'float': 'DOUBLE PRECISION',
        'string': 'TEXT',
        'array': 'JSONB',
        'object': 'JSONB',
        'date': 'TIMESTAMP',
        'objectid': 'TEXT',
    }

    def __init__(self, sample_size: int = 100):
        """
        Initialize schema generator.
        
        Args:
            sample_size: Number of documents to sample for type inference.
        """
        self.sample_size = sample_size
        self.field_types: Dict[str, Set[str]] = {}
        self.final_schema: Dict[str, str] = {}

    def infer_type(self, value: Any) -> str:
        """
        Infer the MongoDB type of a value.
        
        Args:
            value: The value to inspect.
            
        Returns:
            A type string suitable for TYPE_MAPPING.
        """
        if value is None:
            return 'null'
        if isinstance(value, bool):
            return 'boolean'
        if isinstance(value, int) and not isinstance(value, bool):
            return 'integer'
        if isinstance(value, float):
            return 'float'
        if isinstance(value, str):
            return 'string'
        if isinstance(value, list):
            return 'array'
        if isinstance(value, dict):
            return 'object'
        # Handle MongoDB-specific types
        if hasattr(value, '__class__'):
            type_name = value.__class__.__name__.lower()
            if 'objectid' in type_name:
                return 'objectid'
            if 'datetime' in type_name:
                return 'date'
        return 'string'

    def merge_types(self, types: Set[str]) -> str:
        """
        Merge multiple observed types for a field into a single PostgreSQL type.
        Priority: object/array > string > float > integer > boolean > null
        
        Args:
            types: Set of type strings.
            
        Returns:
            The merged PostgreSQL type.
        """
        if not types:
            return self.TYPE_MAPPING['null']
        
        # Priority order
        priority = ['array', 'object', 'string', 'float', 'integer', 'boolean', 'null', 'date']
        for ptype in priority:
            if ptype in types:
                return self.TYPE_MAPPING.get(ptype, self.TYPE_MAPPING['string'])
        
        return self.TYPE_MAPPING['string']

    def analyze_collection(
        self, 
        mongo_db: Any, 
        collection_name: Any
    ) -> Dict[str, str]:
        """
        Analyze a MongoDB collection and infer field types.
        
        Args:
            mongo_db: A pymongo database object or None if collection_name is a collection.
            collection_name: Name of the collection (str) or a pymongo Collection object directly.
            
        Returns:
            Dictionary mapping field names to PostgreSQL types.
        """
        self.field_types = {}
        # Handle both cases: collection name string or collection object
        if isinstance(collection_name, str):
            collection = mongo_db[collection_name]
        else:
            # Assume it's already a pymongo Collection object
            collection = collection_name
        
        # Sample documents from the collection
        sample = list(collection.find({}).limit(self.sample_size))
        
        if not sample:
            logger.warning(f"Collection '{collection_name}' is empty; no schema inferred.")
            return {}
        
        logger.info(f"Analyzing {len(sample)} documents from '{collection_name}'")
        
        # Gather field types
        for doc in sample:
            for field_name, value in doc.items():
                inferred = self.infer_type(value)
                self.field_types.setdefault(field_name, set()).add(inferred)
        
        # Merge types
        self.final_schema = {}
        for field_name, types in self.field_types.items():
            merged_type = self.merge_types(types)
            self.final_schema[field_name] = merged_type
            logger.debug(f"  {field_name}: {merged_type} (observed: {types})")
        
        return self.final_schema

    def generate_create_table_sql(
        self,
        table_name: str,
        schema_name: str = 'public',
        pk_field: Optional[str] = None,
        exclude_fields: Optional[List[str]] = None,
    ) -> str:
        """
        Generate a CREATE TABLE statement from the inferred schema.
        
        Args:
            table_name: Name of the table to create.
            schema_name: PostgreSQL schema (default: 'public').
            pk_field: Field to use as primary key (e.g., 'id', '_id', 'seriesid').
            exclude_fields: List of field names to exclude.
            
        Returns:
            A CREATE TABLE SQL statement.
        """
        if not self.final_schema:
            raise ValueError("No schema has been analyzed. Call analyze_collection first.")
        
        exclude_fields = exclude_fields or []
        columns = []
        
        for field_name, pg_type in self.final_schema.items():
            if field_name in exclude_fields or field_name == '_id':
                continue
            columns.append(f"  {field_name} {pg_type}")
        
        # Add primary key constraint if specified
        constraints = []
        if pk_field and pk_field in self.final_schema:
            constraints.append(f"  PRIMARY KEY ({pk_field})")
        
        all_lines = columns + constraints
        columns_str = ',\n'.join(all_lines)
        
        full_table_name = f"{schema_name}.{table_name}"
        sql = f"""CREATE TABLE IF NOT EXISTS {full_table_name} (
{columns_str}
);"""
        
        return sql

    def create_table_in_postgres(
        self,
        engine: Any,
        table_name: str,
        schema_name: str = 'public',
        pk_field: Optional[str] = None,
        exclude_fields: Optional[List[str]] = None,
        drop_existing: bool = False,
    ) -> None:
        """
        Create a table in PostgreSQL based on the inferred schema.
        
        Args:
            engine: SQLAlchemy engine connected to PostgreSQL, or a dict with 'engine' key.
            table_name: Name of the table to create.
            schema_name: PostgreSQL schema (default: 'public').
            pk_field: Field to use as primary key.
            exclude_fields: List of field names to exclude.
            drop_existing: If True, drop the table before creating.
        """
        # Handle both dict (from db/functions.py) and engine object directly
        if isinstance(engine, dict):
            engine = engine['engine']
        
        full_table_name = f"{schema_name}.{table_name}"
        
        create_sql = self.generate_create_table_sql(
            table_name,
            schema_name=schema_name,
            pk_field=pk_field,
            exclude_fields=exclude_fields,
        )
        
        with engine.begin() as conn:
            if drop_existing:
                logger.info(f"Dropping existing table {full_table_name}...")
                conn.execute(text(f"DROP TABLE IF EXISTS {full_table_name};"))
            
            logger.info(f"Creating table {full_table_name}...")
            conn.execute(text(create_sql))
            logger.info(f"Table {full_table_name} created successfully.")

    def create_table_from_collection(
        self,
        mongo_db: Any,
        collection_name: str,
        engine: Any,
        table_name: Optional[str] = None,
        schema_name: str = 'dbo',
        pk_field: Optional[str] = None,
        exclude_fields: Optional[List[str]] = None,
        drop_existing: bool = False,
    ) -> str:
        """
        End-to-end: analyze a MongoDB collection and create a PostgreSQL table.
        
        Args:
            mongo_db: A pymongo database object.
            collection_name: Name of the MongoDB collection.
            engine: SQLAlchemy engine connected to PostgreSQL.
            table_name: Name of the PostgreSQL table (defaults to collection_name).
            schema_name: PostgreSQL schema (default: 'dbo').
            pk_field: Field to use as primary key.
            exclude_fields: List of field names to exclude.
            drop_existing: If True, drop the table before creating.
            
        Returns:
            The generated SQL statement.
        """
        if table_name is None:
            table_name = collection_name
        
        # Analyze the collection
        self.analyze_collection(mongo_db, collection_name)
        
        # Create the table
        self.create_table_in_postgres(
            engine,
            table_name,
            schema_name=schema_name,
            pk_field=pk_field,
            exclude_fields=exclude_fields,
            drop_existing=drop_existing,
        )
        
        # Return the SQL for reference
        return self.generate_create_table_sql(
            table_name,
            schema_name=schema_name,
            pk_field=pk_field,
            exclude_fields=exclude_fields,
        )


# Convenience function
def create_postgres_table_from_mongo(
    mongo_db: Any,
    collection_name: str,
    engine: Any,
    table_name: Optional[str] = None,
    schema_name: str = 'dbo',
    pk_field: Optional[str] = None,
    sample_size: int = 100,
    drop_existing: bool = False,
) -> str:
    """
    Convenience function to create a PostgreSQL table from a MongoDB collection.
    
    Example:
        from db.schema_generator import create_postgres_table_from_mongo
        
        sql = create_postgres_table_from_mongo(
            mongo_db=mgdb,
            collection_name='series',
            engine=pg_engine,
            table_name='seriesdata',
            schema_name='dbo',
            pk_field='seriesid',
            drop_existing=True,
        )
        print(f"Created table with schema:\\n{sql}")
    
    Args:
        mongo_db: A pymongo database object.
        collection_name: Name of the MongoDB collection.
        engine: SQLAlchemy engine connected to PostgreSQL.
        table_name: Name of the PostgreSQL table (defaults to collection_name).
        schema_name: PostgreSQL schema (default: 'dbo').
        pk_field: Field to use as primary key (optional).
        sample_size: Number of documents to sample (default: 100).
        drop_existing: If True, drop the table before creating.
        
    Returns:
        The generated SQL statement.
    """
    generator = MongoToPostgresSchemaGenerator(sample_size=sample_size)
    return generator.create_table_from_collection(
        mongo_db=mongo_db,
        collection_name=collection_name,
        engine=engine,
        table_name=table_name,
        schema_name=schema_name,
        pk_field=pk_field,
        drop_existing=drop_existing,
    )


class MongoDocumentInserter:
    """Insert MongoDB documents into PostgreSQL tables with type conversion."""

    def __init__(self, batch_size: int = 10000):
        """
        Initialize the document inserter.
        
        Args:
            batch_size: Number of documents to insert per batch (default: 1000).
        """
        self.batch_size = batch_size

    def convert_value(self, value: Any) -> Any:
        """
        Convert a MongoDB value to PostgreSQL-compatible type.
        
        - MongoDB ObjectId → string
        - datetime → string (ISO format)
        - list/dict → JSON string (will be stored as JSONB)
        - None → None
        
        Args:
            value: The value to convert.
            
        Returns:
            The converted value.
        """
        if value is None:
            return None
        
        # Handle MongoDB ObjectId
        if hasattr(value, '__class__') and 'ObjectId' in value.__class__.__name__:
            return str(value)
        
        # Handle datetime
        if hasattr(value, 'isoformat'):
            return value.isoformat()
        
        # Handle lists and dicts (will be stored as JSONB)
        if isinstance(value, (list, dict)):
            return json.dumps(value)
        
        return value

    def prepare_documents(
        self,
        documents: List[Dict[str, Any]],
        column_names: List[str],
        exclude_fields: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Prepare documents for insertion: convert types, exclude fields, etc.
        
        Args:
            documents: List of documents (dicts) from MongoDB.
            column_names: Column names in the target table.
            exclude_fields: Fields to exclude (e.g., ['_id']).
            
        Returns:
            List of prepared documents ready for insertion.
        """
        exclude_fields = exclude_fields or ['_id']
        prepared = []
        
        for doc in documents:
            row = {}
            for col in column_names:
                if col not in exclude_fields and col in doc:
                    row[col] = self.convert_value(doc[col])
                elif col not in exclude_fields:
                    row[col] = None
            prepared.append(row)
        
        return prepared

    def insert_documents(
        self,
        engine: Any,
        table_name: str,
        documents: List[Dict[str, Any]],
        schema_name: str = 'dbo',
        on_conflict: Optional[str] = None,
        exclude_fields: Optional[List[str]] = None,
        skip_null: bool = False,
    ) -> int:
        """
        Insert documents into a PostgreSQL table in batches.
        
        Args:
            engine: SQLAlchemy engine (or dict with 'engine' key).
            table_name: Target table name.
            documents: List of documents (dicts) to insert.
            schema_name: PostgreSQL schema (default: 'dbo').
            on_conflict: ON CONFLICT clause (e.g., "DO NOTHING" or 
                         "DO UPDATE SET field=EXCLUDED.field").
            exclude_fields: Fields to exclude from insert (default: ['_id']).
            skip_null: If True, skip None values in the INSERT (uses COALESCE).
            
        Returns:
            Total number of documents inserted.
        """
        if isinstance(engine, dict):
            engine = engine['engine']
        
        if not documents:
            logger.warning("No documents to insert.")
            return 0
        
        exclude_fields = exclude_fields or ['_id']
        
        full_table_name = f"{schema_name}.{table_name}"
        total_inserted = 0

        # Reflect target table to get column definitions and types
        metadata = MetaData()
        try:
            table = Table(table_name, metadata, autoload_with=engine, schema=schema_name)
            column_objs = {c.name: c for c in table.columns}
            column_names = list(column_objs.keys())
        except Exception:
            # Fallback: derive column names from first document
            column_names = [k for k in documents[0].keys() if k not in exclude_fields]
            column_objs = {}

        # Prepare and convert documents according to target column types
        prepared: List[Dict[str, Any]] = []
        for doc in documents:
            # create case-insensitive key map for doc
            key_map = {k.lower(): k for k in doc.keys()}
            row: Dict[str, Any] = {}
            for col in column_names:
                if exclude_fields and col in exclude_fields:
                    continue
                # find matching key in document (case-insensitive)
                val = None
                if col in doc:
                    val = doc[col]
                elif col.lower() in key_map:
                    val = doc[key_map[col.lower()]]

                # Convert based on column type if available
                colobj = column_objs.get(col)
                if colobj is not None and val is not None:
                    try:
                        from datetime import datetime
                        import sqlalchemy as sa

                        ctype = colobj.type
                        # Integer target: convert ISO datetime strings to epoch
                        if isinstance(ctype, (sa.Integer, sa.BigInteger)):
                            if isinstance(val, str):
                                # try parse datetime string
                                try:
                                    dt = datetime.fromisoformat(val)
                                    val = int(dt.timestamp())
                                except Exception:
                                    # try numeric parse
                                    try:
                                        val = int(val)
                                    except Exception:
                                        pass
                            elif isinstance(val, float):
                                val = int(val)
                        # Datetime target: convert epoch ints to datetime
                        elif isinstance(ctype, (sa.DateTime, sa.TIMESTAMP)):
                            if isinstance(val, (int, float)):
                                val = datetime.fromtimestamp(val)
                            elif isinstance(val, str):
                                try:
                                    val = datetime.fromisoformat(val)
                                except Exception:
                                    pass
                        # JSON target: ensure Python dict/list
                        elif 'JSON' in type(ctype).__name__.upper() or 'JSON' in str(ctype).upper():
                            if isinstance(val, str):
                                try:
                                    val = json.loads(val)
                                except Exception:
                                    pass
                    except Exception:
                        # ignore conversion errors and use original value
                        pass

                # Final conversion for general values
                if val is None:
                    row[col] = None
                else:
                    row[col] = self.convert_value(val) if col not in (exclude_fields or []) else None

            prepared.append(row)

        # Build INSERT statement template using reflected column order
        cols_str = ', '.join(column_names)
        placeholders = ', '.join([':' + col for col in column_names])
        insert_sql = f"INSERT INTO {full_table_name} ({cols_str}) VALUES ({placeholders})"
        if on_conflict:
            insert_sql += f" {on_conflict}"

        # Use different strategies depending on whether ON CONFLICT is needed.
        # For ON CONFLICT updates, per-row execution is slow; use psycopg2's
        # execute_values to perform fast multi-row INSERT ... VALUES (...) ON CONFLICT ...
        if on_conflict:
            try:
                # Prepare ordered tuples for insertion
                values_list = []
                for doc in prepared:
                    values_list.append(tuple(doc.get(col) for col in column_names))

                # Build base INSERT with placeholder for execute_values
                insert_base = f"INSERT INTO {full_table_name} ({cols_str}) VALUES %s {on_conflict}"

                # Use raw DB-API connection for execute_values
                try:
                    import psycopg2
                except Exception:
                    # Fall back to SQLAlchemy executemany if psycopg2 not available
                    with engine.begin() as conn:
                        for doc in prepared:
                            try:
                                conn.execute(text(insert_sql), [doc])
                                total_inserted += 1
                            except Exception as e:
                                logger.error(f"Error inserting document into {full_table_name}: {e}")
                                raise
                        logger.info(f"Inserted {total_inserted} documents into {full_table_name}")
                else:
                    # Use a single staging temporary table per insert_documents call.
                    # Stream each batch into the temp table via COPY, then run one
                    # INSERT ... SELECT ... ON CONFLICT ... to upsert all rows.
                    raw_conn = engine.raw_connection()
                    try:
                        cur = raw_conn.cursor()
                        import io
                        import csv
                        import time
                        import os

                        temp_name = f"tmp_{table_name}_{int(time.time())}_{os.getpid()}"
                        try:
                            # Create temporary table once
                            cur.execute(f"CREATE TEMP TABLE {temp_name} (LIKE {full_table_name} INCLUDING ALL);")
                            # Stream batches into temp table
                            for i in range(0, len(prepared), self.batch_size):
                                batch = prepared[i:i + self.batch_size]
                                try:
                                    sio = io.StringIO()
                                    writer = csv.writer(sio)
                                    for doc in batch:
                                        row = []
                                        for col in column_names:
                                            val = doc.get(col)
                                            if val is None:
                                                row.append('\\N')
                                            else:
                                                row.append(val)
                                        writer.writerow(row)
                                    sio.seek(0)

                                    copy_sql = f"COPY {temp_name} ({cols_str}) FROM STDIN WITH CSV NULL '\\N'"
                                    cur.copy_expert(copy_sql, sio)
                                    raw_conn.commit()
                                    batch_count = len(batch)
                                    total_inserted += batch_count
                                    logger.info(f"Copied {batch_count} rows into staging {temp_name} (total staged: {total_inserted})")
                                except Exception as e:
                                    raw_conn.rollback()
                                    logger.error(f"Error copying batch into {temp_name}: {e}")
                                    raise

                            # Perform a single upsert from the staging table
                            try:
                                insert_from_temp = f"INSERT INTO {full_table_name} ({cols_str}) SELECT {cols_str} FROM {temp_name} {on_conflict}"
                                cur.execute(insert_from_temp)
                                raw_conn.commit()
                                logger.info(f"Upserted staged rows from {temp_name} into {full_table_name}")
                            except Exception as e:
                                raw_conn.rollback()
                                logger.error(f"Error upserting from {temp_name} into {full_table_name}: {e}")
                                raise
                        finally:
                            # Temp table will be dropped on commit/connection close, ensure commit
                            try:
                                raw_conn.commit()
                            except Exception:
                                pass
                    finally:
                        try:
                            cur.close()
                        except Exception:
                            pass
                        try:
                            raw_conn.close()
                        except Exception:
                            pass
            except Exception:
                raise
        else:
            with engine.begin() as conn:
                for i in range(0, len(prepared), self.batch_size):
                    batch = prepared[i:i + self.batch_size]
                    try:
                        conn.execute(text(insert_sql), batch)
                        batch_count = len(batch)
                        total_inserted += batch_count
                        logger.info(
                            f"Inserted {batch_count} documents into {full_table_name} (total: {total_inserted})"
                        )
                    except Exception as e:
                        logger.error(f"Error inserting batch into {full_table_name}: {e}")
                        raise

        return total_inserted

    def insert_from_collection(
        self,
        engine: Any,
        table_name: str,
        collection: Any,
        schema_name: str = 'dbo',
        on_conflict: Optional[str] = None,
        exclude_fields: Optional[List[str]] = None,
        query_filter: Optional[Dict] = None,
        limit: Optional[int] = None,
    ) -> int:
        """
        Insert documents from a MongoDB collection into PostgreSQL.
        
        Args:
            engine: SQLAlchemy engine (or dict with 'engine' key).
            table_name: Target PostgreSQL table name.
            collection: PyMongo collection object.
            schema_name: PostgreSQL schema (default: 'dbo').
            on_conflict: ON CONFLICT clause.
            exclude_fields: Fields to exclude (default: ['_id']).
            query_filter: MongoDB query filter (default: {}).
            limit: Maximum documents to insert (default: None for all).
            
        Returns:
            Total number of documents inserted.
        """
        query_filter = query_filter or {}
        
        logger.info(f"Fetching documents from MongoDB collection...")
        cursor = collection.find(query_filter)
        
        if limit:
            cursor = cursor.limit(limit)
        
        documents = list(cursor)
        logger.info(f"Retrieved {len(documents)} documents.")
        
        return self.insert_documents(
            engine=engine,
            table_name=table_name,
            documents=documents,
            schema_name=schema_name,
            on_conflict=on_conflict,
            exclude_fields=exclude_fields,
        )


# Convenience function for inserting documents
def insert_mongo_documents_to_postgres(
    engine: Any,
    table_name: str,
    documents: List[Dict[str, Any]],
    schema_name: str = 'dbo',
    on_conflict: Optional[str] = None,
    exclude_fields: Optional[List[str]] = None,
    batch_size: int = 1000,
) -> int:
    """
    Convenience function to insert documents into PostgreSQL.
    
    Example:
        from db.schema_generator import insert_mongo_documents_to_postgres
        
        count = insert_mongo_documents_to_postgres(
            engine=dbengine,
            table_name='seriesdata',
            documents=series_docs,
            schema_name='dbo',
            on_conflict="DO NOTHING",
            exclude_fields=['_id'],
        )
        print(f"Inserted {count} documents")
    
    Args:
        engine: SQLAlchemy engine (or dict with 'engine' key).
        table_name: Target table name.
        documents: List of documents to insert.
        schema_name: PostgreSQL schema (default: 'dbo').
        on_conflict: ON CONFLICT clause (optional).
        exclude_fields: Fields to exclude (default: ['_id']).
        batch_size: Number of documents per batch (default: 1000).
        
    Returns:
        Total number of documents inserted.
    """
    inserter = MongoDocumentInserter(batch_size=batch_size)
    return inserter.insert_documents(
        engine=engine,
        table_name=table_name,
        documents=documents,
        schema_name=schema_name,
        on_conflict=on_conflict,
        exclude_fields=exclude_fields,
    )

