"""MongoDB sink using pymongo."""

from __future__ import annotations

import logging
from typing import Any

from .base import BaseSink

logger = logging.getLogger(__name__)

try:
    import pymongo  # type: ignore[import-untyped]
    from pymongo import UpdateOne  # type: ignore[import-untyped]
    from pymongo.errors import BulkWriteError  # type: ignore[import-untyped]
except ImportError:
    pymongo = None  # type: ignore[assignment]
    UpdateOne = None  # type: ignore[assignment]
    BulkWriteError = None  # type: ignore[assignment]


def _to_doc(record: dict[str, Any]) -> dict[str, Any]:
    doc = dict(record)
    if "id" in doc:
        doc["_id"] = doc.pop("id")
    return doc


class MongoDBSink(BaseSink):
    def __init__(
        self,
        uri: str = "mongodb://localhost:27017",
        database: str = "extractor",
        collection: str = "elements",
        batch_size: int = 500,
        upsert: bool = False,
    ) -> None:
        if pymongo is None:
            raise ImportError("Run: pip install 'coresdk-extractor[mongodb]'")
        self.uri = uri
        self.database = database
        self.collection = collection
        self.batch_size = batch_size
        self.upsert = upsert
        self._client: Any = None
        self._col: Any = None
        self._buffer: list[dict[str, Any]] = []

    def connect(self) -> None:
        self._client = pymongo.MongoClient(self.uri)
        db = self._client[self.database]
        self._col = db[self.collection]
        self._col.create_index("source_sha256")
        self._col.create_index("element_type")
        self._col.create_index(
            [("source_filename", pymongo.ASCENDING), ("sequence_index", pymongo.ASCENDING)]
        )

    def feed(self, record: dict[str, Any]) -> None:
        self._buffer.append(record)
        if len(self._buffer) >= self.batch_size:
            self.write_batch(self._buffer)
            self._buffer = []

    def write_batch(self, records: list[dict[str, Any]]) -> None:
        if not records or self._col is None:
            return
        docs = [_to_doc(r) for r in records]
        if self.upsert:
            ops = [UpdateOne({"_id": doc["_id"]}, {"$set": doc}, upsert=True) for doc in docs]
            try:
                self._col.bulk_write(ops, ordered=False)
            except BulkWriteError as exc:
                n_errors = len(exc.details.get("writeErrors", []))
                logger.warning(
                    "MongoDB bulk_write completed with %d write error(s): %s", n_errors, exc.details
                )
            except Exception as exc:
                logger.error("MongoDB bulk_write failed: %s", exc)
                raise
        else:
            try:
                self._col.insert_many(docs, ordered=False)
            except Exception as e:
                # Gracefully handle duplicate key errors on re-run
                if (
                    BulkWriteError is not None
                    and isinstance(e, BulkWriteError)
                    or "duplicate key" in str(e).lower()
                    or "E11000" in str(e)
                ):
                    pass
                else:
                    raise

    def flush(self) -> None:
        if self._buffer:
            self.write_batch(self._buffer)
            self._buffer = []

    def close(self) -> None:
        if self._client is not None:
            self._client.close()
            self._client = None
