"""
MongoDB loader for ETL orchestrator.

Supports loading data to MongoDB collections using:
- insert: Insert new documents (fails on duplicates)
- upsert: Insert or update based on primary key
- replace: Replace entire documents
- update: Update specific fields in existing documents
- delete: Delete documents by primary key

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

    write_method: upsert  # insert, upsert, replace, update, delete
    primary_key: _id      # field to use for matching (default: _id)

    # Optional: ordered bulk writes (default: true)
    # Set to false for better performance when order doesn't matter
    ordered: false

    # Optional: bypass document validation
    bypass_validation: false
    ```
"""

from __future__ import annotations

import logging
import time
from typing import Any

from bson import ObjectId
from pymongo import MongoClient
from pymongo.operations import DeleteOne, InsertOne, ReplaceOne, UpdateOne

from pycharter.etl_generator.database import (
    DEFAULT_TUNNEL_LOCAL_PORT,
    create_ssh_tunnel,
)
from pycharter.etl_generator.loaders.base import BaseLoader
from pycharter.etl_generator.result import LoadResult
from pycharter.utils.value_injector import resolve_values

logger = logging.getLogger(__name__)

# Default MongoDB port
DEFAULT_MONGODB_PORT = 27017


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

    pattern = r"mongodb(\+srv)?://([^@]+@)?([^/:]+)(:\d+)?"

    def replace_host(match):
        protocol = match.group(1) or ""
        auth = match.group(2) or ""
        return f"mongodb://{auth}127.0.0.1:{local_port}"

    return re.sub(pattern, replace_host, url)


def _convert_id_field(doc: dict[str, Any], primary_key: str) -> dict[str, Any]:
    """
    Convert primary key to ObjectId if it's a valid ObjectId string.

    Args:
        doc: Document dictionary
        primary_key: Primary key field name

    Returns:
        Document with converted _id if applicable
    """
    result = doc.copy()

    # If primary key is '_id' and value is a valid ObjectId string, convert it
    if primary_key == "_id" and "_id" in result:
        id_value = result["_id"]
        if isinstance(id_value, str) and ObjectId.is_valid(id_value):
            result["_id"] = ObjectId(id_value)

    return result


class MongoDBLoader(BaseLoader):
    """
    Loader for MongoDB data targets.

    Supports:
    - Insert (insert_many for bulk inserts)
    - Upsert (bulk_write with ReplaceOne operations)
    - Replace (bulk_write with ReplaceOne)
    - Update (bulk_write with UpdateOne)
    - Delete (bulk_write with DeleteOne)

    Example (programmatic API):
        >>> loader = MongoDBLoader(
        ...     connection_string="mongodb://localhost:27017",
        ...     database="mydb",
        ...     collection="users",
        ...     write_method="upsert",
        ...     primary_key="_id",
        ... )
        >>> result = await loader.load(data)

    Example (config-driven):
        >>> loader = MongoDBLoader.from_config(config)
        >>> result = await loader.load(data)
    """

    def __init__(
        self,
        connection_string: str,
        database: str,
        collection: str,
        write_method: str = "upsert",
        primary_key: str | list[str] = "_id",
        batch_size: int = 1000,
        ordered: bool = True,
        bypass_validation: bool = False,
        ssh_tunnel: dict[str, Any] | None = None,
    ):
        """
        Initialize MongoDB loader.

        Args:
            connection_string: MongoDB connection URL
            database: Database name
            collection: Collection name
            write_method: Write method (insert, upsert, replace, update, delete)
            primary_key: Field(s) to use for matching documents
            batch_size: Number of operations per bulk write
            ordered: Whether bulk writes should be ordered
            bypass_validation: Whether to bypass document validation
            ssh_tunnel: SSH tunnel configuration
        """
        self.connection_string = connection_string
        self.database = database
        self.collection = collection
        self.write_method = write_method
        self.primary_key = primary_key
        self.batch_size = batch_size
        self.ordered = ordered
        self.bypass_validation = bypass_validation
        self.ssh_tunnel = ssh_tunnel

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "MongoDBLoader":
        """Create loader 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"),
            write_method=config.get("write_method", "upsert"),
            primary_key=config.get("primary_key", "_id"),
            batch_size=config.get("batch_size", 1000),
            ordered=config.get("ordered", True),
            bypass_validation=config.get("bypass_validation", False),
            ssh_tunnel=mongo_config.get("ssh_tunnel"),
        )

    def _build_filter(self, doc: dict[str, Any]) -> dict[str, Any]:
        """
        Build a filter document for matching.

        Args:
            doc: Source document

        Returns:
            Filter document using primary key field(s)
        """
        if isinstance(self.primary_key, list):
            # Composite key
            return {key: doc.get(key) for key in self.primary_key if key in doc}
        else:
            # Single key
            key_value = doc.get(self.primary_key)
            if key_value is None:
                raise ValueError(
                    f"Document missing primary key field '{self.primary_key}': {doc}"
                )
            return {self.primary_key: key_value}

    async def load(self, data: list[dict[str, Any]], **params) -> LoadResult:
        """
        Load data to MongoDB collection.

        Args:
            data: List of documents to load
            **params: Additional parameters (config_context, contract_dir)

        Returns:
            LoadResult with statistics
        """
        start_time = time.time()

        if not data:
            return LoadResult(success=True, rows_loaded=0)

        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")

        # Handle SSH tunnel if configured
        tunnel = None
        connection_string = self.connection_string

        config_context = params.get("config_context")
        contract_dir = params.get("contract_dir")
        source_file = str(contract_dir / "load.yaml") if contract_dir else None

        if self.ssh_tunnel:
            ssh_config = resolve_values(
                self.ssh_tunnel, 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):
                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)
                    )
                    connection_string = _modify_mongodb_url_for_tunnel(
                        connection_string, local_port
                    )

        # Create MongoDB client
        client = MongoClient(connection_string)
        db = client[self.database]
        collection = db[self.collection]

        inserted = 0
        updated = 0
        deleted = 0
        errors = []

        try:
            logger.info(
                "Loading to MongoDB: %s.%s (method: %s, records: %s)",
                self.database,
                self.collection,
                self.write_method,
                len(data),
            )

            # Process in batches
            for i in range(0, len(data), self.batch_size):
                batch = data[i : i + self.batch_size]

                if self.write_method == "insert":
                    # Simple insert_many
                    result = collection.insert_many(
                        batch,
                        ordered=self.ordered,
                        bypass_document_validation=self.bypass_validation,
                    )
                    inserted += len(result.inserted_ids)

                elif self.write_method == "upsert":
                    # Bulk upsert using ReplaceOne with upsert=True
                    operations = []
                    for doc in batch:
                        doc = _convert_id_field(
                            doc,
                            (
                                self.primary_key
                                if isinstance(self.primary_key, str)
                                else "_id"
                            ),
                        )
                        filter_doc = self._build_filter(doc)
                        operations.append(ReplaceOne(filter_doc, doc, upsert=True))

                    if operations:
                        result = collection.bulk_write(
                            operations,
                            ordered=self.ordered,
                            bypass_document_validation=self.bypass_validation,
                        )
                        inserted += result.upserted_count
                        updated += result.modified_count

                elif self.write_method == "replace":
                    # Bulk replace (no upsert - only updates existing)
                    operations = []
                    for doc in batch:
                        doc = _convert_id_field(
                            doc,
                            (
                                self.primary_key
                                if isinstance(self.primary_key, str)
                                else "_id"
                            ),
                        )
                        filter_doc = self._build_filter(doc)
                        operations.append(ReplaceOne(filter_doc, doc, upsert=False))

                    if operations:
                        result = collection.bulk_write(
                            operations,
                            ordered=self.ordered,
                            bypass_document_validation=self.bypass_validation,
                        )
                        updated += result.modified_count

                elif self.write_method == "update":
                    # Bulk update using $set
                    operations = []
                    for doc in batch:
                        doc = _convert_id_field(
                            doc,
                            (
                                self.primary_key
                                if isinstance(self.primary_key, str)
                                else "_id"
                            ),
                        )
                        filter_doc = self._build_filter(doc)
                        # Remove primary key from update fields
                        update_fields = {
                            k: v
                            for k, v in doc.items()
                            if k != self.primary_key
                            and k
                            not in (
                                self.primary_key
                                if isinstance(self.primary_key, list)
                                else [self.primary_key]
                            )
                        }
                        operations.append(
                            UpdateOne(filter_doc, {"$set": update_fields})
                        )

                    if operations:
                        result = collection.bulk_write(
                            operations,
                            ordered=self.ordered,
                            bypass_document_validation=self.bypass_validation,
                        )
                        updated += result.modified_count

                elif self.write_method == "delete":
                    # Bulk delete by primary key
                    operations = []
                    for doc in batch:
                        doc = _convert_id_field(
                            doc,
                            (
                                self.primary_key
                                if isinstance(self.primary_key, str)
                                else "_id"
                            ),
                        )
                        filter_doc = self._build_filter(doc)
                        operations.append(DeleteOne(filter_doc))

                    if operations:
                        result = collection.bulk_write(
                            operations,
                            ordered=self.ordered,
                        )
                        deleted += result.deleted_count

                elif self.write_method == "truncate_and_load":
                    # Delete all documents first (only on first batch)
                    if i == 0:
                        collection.delete_many({})
                        logger.info(
                            "Truncated collection %s.%s", self.database, self.collection
                        )

                    # Then insert
                    result = collection.insert_many(
                        batch,
                        ordered=self.ordered,
                        bypass_document_validation=self.bypass_validation,
                    )
                    inserted += len(result.inserted_ids)

                else:
                    raise ValueError(f"Unsupported write method: {self.write_method}")

            duration = time.time() - start_time
            total_affected = inserted + updated + deleted

            logger.info(
                "MongoDB load completed: %s records affected "
                "(inserted: %s, updated: %s, deleted: %s) in %.2fs",
                total_affected,
                inserted,
                updated,
                deleted,
                duration,
            )

            return LoadResult(
                success=True,
                rows_loaded=inserted + updated,
                duration_seconds=duration,
                metadata={
                    "inserted": inserted,
                    "updated": updated,
                    "deleted": deleted,
                },
            )

        except Exception as e:
            logger.error("MongoDB load failed: %s", e, exc_info=True)
            return LoadResult(
                success=False,
                error=str(e),
                duration_seconds=time.time() - start_time,
            )
        finally:
            client.close()
            if tunnel:
                tunnel.stop()
