"""
MongoDB extractor for ETL orchestrator.

Supports extracting data from MongoDB collections using:
- Simple find() queries with filters
- Aggregation pipelines for complex transformations

Configuration example:
    ```yaml
    type: mongodb
    mongodb:
      url: mongodb://user:pass@localhost:27017
      database: mydb
      collection: users
      ssh_tunnel:
        enabled: false
        host: bastion.example.com
        username: tunnel_user
        remote_host: mongo-internal.example.com
        remote_port: 27017

    # Option 1: Simple query filter
    query:
      status: active
      created_at:
        $gte: "2024-01-01"

    # Option 2: Aggregation pipeline (takes precedence over query)
    pipeline:
      - $match:
          status: active
      - $project:
          name: 1
          email: 1
          _id: 0

    # Optional: projection for simple queries
    projection:
      name: 1
      email: 1
      _id: 0

    # Optional: sorting
    sort:
      created_at: -1

    # Optional: limit results
    limit: 1000
    ```
"""

from __future__ import annotations

import logging
from typing import Any, AsyncIterator

from bson import ObjectId
from pymongo import MongoClient

from pycharter.etl_generator.database import (
    DEFAULT_TUNNEL_LOCAL_PORT,
    create_ssh_tunnel,
)
from pycharter.etl_generator.extractors.base import BaseExtractor
from pycharter.utils.value_injector import resolve_values

logger = logging.getLogger(__name__)

# Default MongoDB port
DEFAULT_MONGODB_PORT = 27017


def _serialize_document(doc: dict[str, Any]) -> dict[str, Any]:
    """
    Serialize a MongoDB document for JSON compatibility.

    Converts:
    - ObjectId to string
    - Other BSON types as needed

    Args:
        doc: MongoDB document

    Returns:
        JSON-serializable dictionary
    """
    result = {}
    for key, value in doc.items():
        if isinstance(value, ObjectId):
            result[key] = str(value)
        elif isinstance(value, dict):
            result[key] = _serialize_document(value)
        elif isinstance(value, list):
            result[key] = [
                (
                    _serialize_document(item)
                    if isinstance(item, dict)
                    else str(item) if isinstance(item, ObjectId) else item
                )
                for item in value
            ]
        else:
            result[key] = value
    return result


def _modify_mongodb_url_for_tunnel(url: str, local_port: int) -> str:
    """
    Modify MongoDB URL to use local tunnel port.

    Args:
        url: Original MongoDB URL
        local_port: Local tunnel port

    Returns:
        Modified URL pointing to localhost tunnel
    """
    import re

    # Match mongodb:// or mongodb+srv:// URLs
    # Replace host:port with 127.0.0.1:local_port
    # Handle URLs with or without port
    pattern = r"mongodb(\+srv)?://([^@]+@)?([^/:]+)(:\d+)?"

    def replace_host(match):
        protocol = match.group(1) or ""
        auth = match.group(2) or ""
        # For tunneled connections, we don't use SRV
        return f"mongodb://{auth}127.0.0.1:{local_port}"

    return re.sub(pattern, replace_host, url)


class MongoDBExtractor(BaseExtractor):
    """
    Extractor for MongoDB data sources.

    Supports two query modes:
    1. Simple find() queries with optional projection and sorting
    2. Aggregation pipelines for complex data transformations

    Example (programmatic API):
        >>> extractor = MongoDBExtractor(
        ...     connection_string="mongodb://localhost:27017",
        ...     database="mydb",
        ...     collection="users",
        ...     query={"status": "active"},
        ... )
        >>> async for batch in extractor.extract():
        ...     process(batch)

    Example (config-driven):
        >>> extractor = MongoDBExtractor()
        >>> async for batch in extractor.extract_streaming(config, params, headers):
        ...     process(batch)
    """

    def __init__(
        self,
        connection_string: str | None = None,
        database: str | None = None,
        collection: str | None = None,
        query: dict[str, Any] | None = None,
        pipeline: list[dict[str, Any]] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, int] | None = None,
        batch_size: int = 1000,
        max_records: int | None = None,
        ssh_tunnel: dict[str, Any] | None = None,
    ):
        """
        Initialize MongoDB extractor.

        Args:
            connection_string: MongoDB connection URL
            database: Database name
            collection: Collection name
            query: MongoDB query filter (for find operations)
            pipeline: Aggregation pipeline (takes precedence over query)
            projection: Fields to include/exclude (for find operations)
            sort: Sort specification (for find operations)
            batch_size: Number of records per batch
            max_records: Maximum records to extract (None = all)
            ssh_tunnel: SSH tunnel configuration
        """
        self.connection_string = connection_string
        self.database = database
        self.collection = collection
        self.query = query or {}
        self.pipeline = pipeline
        self.projection = projection
        self.sort = sort
        self.batch_size = batch_size
        self.max_records = max_records
        self.ssh_tunnel = ssh_tunnel

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "MongoDBExtractor":
        """Create extractor from configuration dict."""
        mongo_config = config.get("mongodb", {})
        return cls(
            connection_string=mongo_config.get("url")
            or config.get("connection_string"),
            database=mongo_config.get("database") or config.get("database"),
            collection=mongo_config.get("collection") or config.get("collection"),
            query=config.get("query", {}),
            pipeline=config.get("pipeline"),
            projection=config.get("projection"),
            sort=config.get("sort"),
            batch_size=config.get("batch_size", 1000),
            max_records=config.get("max_records"),
            ssh_tunnel=mongo_config.get("ssh_tunnel"),
        )

    async def extract(self, **params) -> AsyncIterator[list[dict[str, Any]]]:
        """
        Extract data from MongoDB.

        Yields:
            Batches of records as dictionaries
        """
        if not self.connection_string:
            raise ValueError("MongoDB connection string is required")
        if not self.database:
            raise ValueError("MongoDB database name is required")
        if not self.collection:
            raise ValueError("MongoDB collection name is required")

        extract_config = {
            "mongodb": {
                "url": self.connection_string,
                "database": self.database,
                "collection": self.collection,
                "ssh_tunnel": self.ssh_tunnel,
            },
            "query": self.query,
            "pipeline": self.pipeline,
            "projection": self.projection,
            "sort": self.sort,
        }

        async for batch in self.extract_streaming(
            extract_config,
            params,
            {},
            batch_size=self.batch_size,
            max_records=self.max_records,
        ):
            yield batch

    def validate_config(self, extract_config: dict[str, Any]) -> None:
        """Validate MongoDB extractor configuration."""
        source_type = extract_config.get("type") or extract_config.get("source_type")
        if source_type and source_type not in ("mongodb", "mongo"):
            raise ValueError(
                f"MongoDBExtractor requires type='mongodb', got '{source_type}'"
            )

        mongo_config = extract_config.get("mongodb", {})
        if not mongo_config.get("url"):
            raise ValueError(
                "MongoDB extractor requires 'mongodb.url' in extract_config"
            )
        if not mongo_config.get("database"):
            raise ValueError(
                "MongoDB extractor requires 'mongodb.database' in extract_config"
            )
        if not mongo_config.get("collection"):
            raise ValueError(
                "MongoDB extractor requires 'mongodb.collection' in extract_config"
            )

    async def extract_streaming(
        self,
        extract_config: dict[str, Any],
        params: dict[str, Any],
        headers: dict[str, Any],
        contract_dir: Any | None = None,
        batch_size: int = 1000,
        max_records: int | None = None,
        config_context: dict[str, Any] | None = None,
    ) -> AsyncIterator[list[dict[str, Any]]]:
        """
        Extract data from MongoDB using find() or aggregation pipeline.

        Yields batches of records as lists of dictionaries.
        """
        # Get MongoDB configuration
        mongo_config = extract_config.get("mongodb", {})

        # Resolve variables
        source_file = str(contract_dir / "extract.yaml") if contract_dir else None
        mongo_url = resolve_values(
            mongo_config.get("url"), context=config_context, source_file=source_file
        )
        db_name = resolve_values(
            mongo_config.get("database"),
            context=config_context,
            source_file=source_file,
        )
        collection_name = resolve_values(
            mongo_config.get("collection"),
            context=config_context,
            source_file=source_file,
        )

        if not mongo_url:
            raise ValueError("MongoDB URL is required")
        if not db_name:
            raise ValueError("MongoDB database name is required")
        if not collection_name:
            raise ValueError("MongoDB collection name is required")

        # Handle SSH tunnel if configured
        ssh_config = mongo_config.get("ssh_tunnel", {})
        tunnel = None
        if ssh_config:
            ssh_config = resolve_values(
                ssh_config, context=config_context, source_file=source_file
            )
            enabled_value = ssh_config.get("enabled", False)
            if isinstance(enabled_value, str):
                enabled_lower = enabled_value.lower()
                ssh_config["enabled"] = enabled_lower in ("true", "1", "yes", "on")
            elif not isinstance(enabled_value, bool):
                ssh_config["enabled"] = bool(enabled_value)

            if ssh_config.get("enabled", False):
                # Set default remote port for MongoDB if not specified
                if "remote_port" not in ssh_config:
                    ssh_config["remote_port"] = DEFAULT_MONGODB_PORT

                tunnel = create_ssh_tunnel(ssh_config)
                if tunnel:
                    local_port = int(
                        ssh_config.get("local_port", DEFAULT_TUNNEL_LOCAL_PORT)
                    )
                    mongo_url = _modify_mongodb_url_for_tunnel(mongo_url, local_port)

        # Get query/pipeline configuration
        query = extract_config.get("query", {})
        pipeline = extract_config.get("pipeline")
        projection = extract_config.get("projection")
        sort_spec = extract_config.get("sort")
        limit = extract_config.get("limit")

        # Merge params into query for parameterized queries
        if params:
            query = {**query, **params}

        # Create MongoDB client
        client = MongoClient(mongo_url)
        db = client[db_name]
        collection = db[collection_name]

        try:
            logger.info(
                "Extracting from MongoDB: %s.%s (pipeline: %s, query: %s)",
                db_name,
                collection_name,
                bool(pipeline),
                bool(query),
            )

            # Choose extraction method
            if pipeline:
                # Use aggregation pipeline
                cursor = collection.aggregate(pipeline, batchSize=batch_size)
            else:
                # Use find() with optional projection, sort, limit
                cursor = collection.find(query, projection)

                if sort_spec:
                    # Convert dict to list of tuples for pymongo
                    sort_list = [(k, v) for k, v in sort_spec.items()]
                    cursor = cursor.sort(sort_list)

                if limit:
                    cursor = cursor.limit(limit)

                cursor = cursor.batch_size(batch_size)

            # Stream results in batches
            current_batch = []
            total_extracted = 0

            for doc in cursor:
                if max_records and total_extracted >= max_records:
                    break

                # Serialize document for JSON compatibility
                record = _serialize_document(doc)
                current_batch.append(record)
                total_extracted += 1

                if len(current_batch) >= batch_size:
                    yield current_batch
                    current_batch = []

            # Yield remaining records
            if current_batch:
                yield current_batch

            logger.info(
                "MongoDB extraction completed: %s records extracted", total_extracted
            )

        except Exception as e:
            logger.error("MongoDB extraction error: %s", e, exc_info=True)
            raise RuntimeError(f"MongoDB extraction failed: {e}") from e
        finally:
            client.close()
            if tunnel:
                tunnel.stop()
