# Copyright 2014-present MongoDB, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""The bulk write operations interface.

.. versionadded:: 2.7
"""
from __future__ import annotations

import copy
import datetime
import logging
from collections.abc import MutableMapping
from itertools import islice
from typing import (
    TYPE_CHECKING,
    Any,
    Iterator,
    Mapping,
    Optional,
    Type,
    Union,
)

from bson.objectid import ObjectId
from bson.raw_bson import RawBSONDocument
from pymongo import _csot, common
from pymongo.asynchronous.client_session import AsyncClientSession, _validate_session_write_concern
from pymongo.asynchronous.helpers import _handle_reauth
from pymongo.bulk_shared import (
    _COMMANDS,
    _DELETE_ALL,
    _merge_command,
    _raise_bulk_write_error,
    _Run,
)
from pymongo.common import (
    validate_is_document_type,
    validate_ok_for_replace,
    validate_ok_for_update,
)
from pymongo.errors import (
    ConfigurationError,
    InvalidOperation,
    NotPrimaryError,
    OperationFailure,
)
from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES
from pymongo.logger import _COMMAND_LOGGER, _CommandStatusMessage, _debug_log
from pymongo.message import (
    _DELETE,
    _INSERT,
    _UPDATE,
    _BulkWriteContext,
    _convert_exception,
    _convert_write_result,
    _EncryptedBulkWriteContext,
    _randint,
)
from pymongo.read_preferences import ReadPreference
from pymongo.write_concern import WriteConcern

if TYPE_CHECKING:
    from pymongo.asynchronous.collection import AsyncCollection
    from pymongo.asynchronous.mongo_client import AsyncMongoClient
    from pymongo.asynchronous.pool import AsyncConnection
    from pymongo.typings import _DocumentOut, _DocumentType, _Pipeline

_IS_SYNC = False


class _AsyncBulk:
    """The private guts of the bulk write API."""

    def __init__(
        self,
        collection: AsyncCollection[_DocumentType],
        ordered: bool,
        bypass_document_validation: bool,
        comment: Optional[str] = None,
        let: Optional[Any] = None,
    ) -> None:
        """Initialize a _AsyncBulk instance."""
        self.collection = collection.with_options(
            codec_options=collection.codec_options._replace(
                unicode_decode_error_handler="replace", document_class=dict
            )
        )
        self.let = let
        if self.let is not None:
            common.validate_is_document_type("let", self.let)
        self.comment: Optional[str] = comment
        self.ordered = ordered
        self.ops: list[tuple[int, Mapping[str, Any]]] = []
        self.executed = False
        self.bypass_doc_val = bypass_document_validation
        self.uses_collation = False
        self.uses_array_filters = False
        self.uses_hint_update = False
        self.uses_hint_delete = False
        self.is_retryable = True
        self.retrying = False
        self.started_retryable_write = False
        # Extra state so that we know where to pick up on a retry attempt.
        self.current_run = None
        self.next_run = None
        self.is_encrypted = False

    @property
    def bulk_ctx_class(self) -> Type[_BulkWriteContext]:
        encrypter = self.collection.database.client._encrypter
        if encrypter and not encrypter._bypass_auto_encryption:
            self.is_encrypted = True
            return _EncryptedBulkWriteContext
        else:
            self.is_encrypted = False
            return _BulkWriteContext

    def add_insert(self, document: _DocumentOut) -> None:
        """Add an insert document to the list of ops."""
        validate_is_document_type("document", document)
        # Generate ObjectId client side.
        if not (isinstance(document, RawBSONDocument) or "_id" in document):
            document["_id"] = ObjectId()
        self.ops.append((_INSERT, document))

    def add_update(
        self,
        selector: Mapping[str, Any],
        update: Union[Mapping[str, Any], _Pipeline],
        multi: bool = False,
        upsert: bool = False,
        collation: Optional[Mapping[str, Any]] = None,
        array_filters: Optional[list[Mapping[str, Any]]] = None,
        hint: Union[str, dict[str, Any], None] = None,
    ) -> None:
        """Create an update document and add it to the list of ops."""
        validate_ok_for_update(update)
        cmd: dict[str, Any] = dict(  # noqa: C406
            [("q", selector), ("u", update), ("multi", multi), ("upsert", upsert)]
        )
        if collation is not None:
            self.uses_collation = True
            cmd["collation"] = collation
        if array_filters is not None:
            self.uses_array_filters = True
            cmd["arrayFilters"] = array_filters
        if hint is not None:
            self.uses_hint_update = True
            cmd["hint"] = hint
        if multi:
            # A bulk_write containing an update_many is not retryable.
            self.is_retryable = False
        self.ops.append((_UPDATE, cmd))

    def add_replace(
        self,
        selector: Mapping[str, Any],
        replacement: Mapping[str, Any],
        upsert: bool = False,
        collation: Optional[Mapping[str, Any]] = None,
        hint: Union[str, dict[str, Any], None] = None,
    ) -> None:
        """Create a replace document and add it to the list of ops."""
        validate_ok_for_replace(replacement)
        cmd = {"q": selector, "u": replacement, "multi": False, "upsert": upsert}
        if collation is not None:
            self.uses_collation = True
            cmd["collation"] = collation
        if hint is not None:
            self.uses_hint_update = True
            cmd["hint"] = hint
        self.ops.append((_UPDATE, cmd))

    def add_delete(
        self,
        selector: Mapping[str, Any],
        limit: int,
        collation: Optional[Mapping[str, Any]] = None,
        hint: Union[str, dict[str, Any], None] = None,
    ) -> None:
        """Create a delete document and add it to the list of ops."""
        cmd = {"q": selector, "limit": limit}
        if collation is not None:
            self.uses_collation = True
            cmd["collation"] = collation
        if hint is not None:
            self.uses_hint_delete = True
            cmd["hint"] = hint
        if limit == _DELETE_ALL:
            # A bulk_write containing a delete_many is not retryable.
            self.is_retryable = False
        self.ops.append((_DELETE, cmd))

    def gen_ordered(self) -> Iterator[Optional[_Run]]:
        """Generate batches of operations, batched by type of
        operation, in the order **provided**.
        """
        run = None
        for idx, (op_type, operation) in enumerate(self.ops):
            if run is None:
                run = _Run(op_type)
            elif run.op_type != op_type:
                yield run
                run = _Run(op_type)
            run.add(idx, operation)
        yield run

    def gen_unordered(self) -> Iterator[_Run]:
        """Generate batches of operations, batched by type of
        operation, in arbitrary order.
        """
        operations = [_Run(_INSERT), _Run(_UPDATE), _Run(_DELETE)]
        for idx, (op_type, operation) in enumerate(self.ops):
            operations[op_type].add(idx, operation)

        for run in operations:
            if run.ops:
                yield run

    @_handle_reauth
    async def write_command(
        self,
        bwc: _BulkWriteContext,
        cmd: MutableMapping[str, Any],
        request_id: int,
        msg: bytes,
        docs: list[Mapping[str, Any]],
        client: AsyncMongoClient,
    ) -> dict[str, Any]:
        """A proxy for SocketInfo.write_command that handles event publishing."""
        cmd[bwc.field] = docs
        if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
            _debug_log(
                _COMMAND_LOGGER,
                clientId=client._topology_settings._topology_id,
                message=_CommandStatusMessage.STARTED,
                command=cmd,
                commandName=next(iter(cmd)),
                databaseName=bwc.db_name,
                requestId=request_id,
                operationId=request_id,
                driverConnectionId=bwc.conn.id,
                serverConnectionId=bwc.conn.server_connection_id,
                serverHost=bwc.conn.address[0],
                serverPort=bwc.conn.address[1],
                serviceId=bwc.conn.service_id,
            )
        if bwc.publish:
            bwc._start(cmd, request_id, docs)
        try:
            reply = await bwc.conn.write_command(request_id, msg, bwc.codec)  # type: ignore[misc]
            duration = datetime.datetime.now() - bwc.start_time
            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
                _debug_log(
                    _COMMAND_LOGGER,
                    clientId=client._topology_settings._topology_id,
                    message=_CommandStatusMessage.SUCCEEDED,
                    durationMS=duration,
                    reply=reply,
                    commandName=next(iter(cmd)),
                    databaseName=bwc.db_name,
                    requestId=request_id,
                    operationId=request_id,
                    driverConnectionId=bwc.conn.id,
                    serverConnectionId=bwc.conn.server_connection_id,
                    serverHost=bwc.conn.address[0],
                    serverPort=bwc.conn.address[1],
                    serviceId=bwc.conn.service_id,
                )
            if bwc.publish:
                bwc._succeed(request_id, reply, duration)  # type: ignore[arg-type]
            await client._process_response(reply, bwc.session)  # type: ignore[arg-type]
        except Exception as exc:
            duration = datetime.datetime.now() - bwc.start_time
            if isinstance(exc, (NotPrimaryError, OperationFailure)):
                failure: _DocumentOut = exc.details  # type: ignore[assignment]
            else:
                failure = _convert_exception(exc)
            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
                _debug_log(
                    _COMMAND_LOGGER,
                    clientId=client._topology_settings._topology_id,
                    message=_CommandStatusMessage.FAILED,
                    durationMS=duration,
                    failure=failure,
                    commandName=next(iter(cmd)),
                    databaseName=bwc.db_name,
                    requestId=request_id,
                    operationId=request_id,
                    driverConnectionId=bwc.conn.id,
                    serverConnectionId=bwc.conn.server_connection_id,
                    serverHost=bwc.conn.address[0],
                    serverPort=bwc.conn.address[1],
                    serviceId=bwc.conn.service_id,
                    isServerSideError=isinstance(exc, OperationFailure),
                )

            if bwc.publish:
                bwc._fail(request_id, failure, duration)
            # Process the response from the server.
            if isinstance(exc, (NotPrimaryError, OperationFailure)):
                await client._process_response(exc.details, bwc.session)  # type: ignore[arg-type]
            raise
        return reply  # type: ignore[return-value]

    async def unack_write(
        self,
        bwc: _BulkWriteContext,
        cmd: MutableMapping[str, Any],
        request_id: int,
        msg: bytes,
        max_doc_size: int,
        docs: list[Mapping[str, Any]],
        client: AsyncMongoClient,
    ) -> Optional[Mapping[str, Any]]:
        """A proxy for AsyncConnection.unack_write that handles event publishing."""
        if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
            _debug_log(
                _COMMAND_LOGGER,
                clientId=client._topology_settings._topology_id,
                message=_CommandStatusMessage.STARTED,
                command=cmd,
                commandName=next(iter(cmd)),
                databaseName=bwc.db_name,
                requestId=request_id,
                operationId=request_id,
                driverConnectionId=bwc.conn.id,
                serverConnectionId=bwc.conn.server_connection_id,
                serverHost=bwc.conn.address[0],
                serverPort=bwc.conn.address[1],
                serviceId=bwc.conn.service_id,
            )
        if bwc.publish:
            cmd = bwc._start(cmd, request_id, docs)
        try:
            result = await bwc.conn.unack_write(msg, max_doc_size)  # type: ignore[func-returns-value, misc, override]
            duration = datetime.datetime.now() - bwc.start_time
            if result is not None:
                reply = _convert_write_result(bwc.name, cmd, result)  # type: ignore[arg-type]
            else:
                # Comply with APM spec.
                reply = {"ok": 1}
                if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
                    _debug_log(
                        _COMMAND_LOGGER,
                        clientId=client._topology_settings._topology_id,
                        message=_CommandStatusMessage.SUCCEEDED,
                        durationMS=duration,
                        reply=reply,
                        commandName=next(iter(cmd)),
                        databaseName=bwc.db_name,
                        requestId=request_id,
                        operationId=request_id,
                        driverConnectionId=bwc.conn.id,
                        serverConnectionId=bwc.conn.server_connection_id,
                        serverHost=bwc.conn.address[0],
                        serverPort=bwc.conn.address[1],
                        serviceId=bwc.conn.service_id,
                    )
            if bwc.publish:
                bwc._succeed(request_id, reply, duration)
        except Exception as exc:
            duration = datetime.datetime.now() - bwc.start_time
            if isinstance(exc, OperationFailure):
                failure: _DocumentOut = _convert_write_result(bwc.name, cmd, exc.details)  # type: ignore[arg-type]
            elif isinstance(exc, NotPrimaryError):
                failure = exc.details  # type: ignore[assignment]
            else:
                failure = _convert_exception(exc)
            if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
                _debug_log(
                    _COMMAND_LOGGER,
                    clientId=client._topology_settings._topology_id,
                    message=_CommandStatusMessage.FAILED,
                    durationMS=duration,
                    failure=failure,
                    commandName=next(iter(cmd)),
                    databaseName=bwc.db_name,
                    requestId=request_id,
                    operationId=request_id,
                    driverConnectionId=bwc.conn.id,
                    serverConnectionId=bwc.conn.server_connection_id,
                    serverHost=bwc.conn.address[0],
                    serverPort=bwc.conn.address[1],
                    serviceId=bwc.conn.service_id,
                    isServerSideError=isinstance(exc, OperationFailure),
                )
            if bwc.publish:
                assert bwc.start_time is not None
                bwc._fail(request_id, failure, duration)
            raise
        return result  # type: ignore[return-value]

    async def _execute_batch_unack(
        self,
        bwc: Union[_BulkWriteContext, _EncryptedBulkWriteContext],
        cmd: dict[str, Any],
        ops: list[Mapping[str, Any]],
        client: AsyncMongoClient,
    ) -> list[Mapping[str, Any]]:
        if self.is_encrypted:
            _, batched_cmd, to_send = bwc.batch_command(cmd, ops)
            await bwc.conn.command(  # type: ignore[misc]
                bwc.db_name,
                batched_cmd,  # type: ignore[arg-type]
                write_concern=WriteConcern(w=0),
                session=bwc.session,  # type: ignore[arg-type]
                client=client,  # type: ignore[arg-type]
            )
        else:
            request_id, msg, to_send = bwc.batch_command(cmd, ops)
            # Though this isn't strictly a "legacy" write, the helper
            # handles publishing commands and sending our message
            # without receiving a result. Send 0 for max_doc_size
            # to disable size checking. Size checking is handled while
            # the documents are encoded to BSON.
            await self.unack_write(bwc, cmd, request_id, msg, 0, to_send, client)  # type: ignore[arg-type]

        return to_send

    async def _execute_batch(
        self,
        bwc: Union[_BulkWriteContext, _EncryptedBulkWriteContext],
        cmd: dict[str, Any],
        ops: list[Mapping[str, Any]],
        client: AsyncMongoClient,
    ) -> tuple[dict[str, Any], list[Mapping[str, Any]]]:
        if self.is_encrypted:
            _, batched_cmd, to_send = bwc.batch_command(cmd, ops)
            result = await bwc.conn.command(  # type: ignore[misc]
                bwc.db_name,
                batched_cmd,  # type: ignore[arg-type]
                codec_options=bwc.codec,
                session=bwc.session,  # type: ignore[arg-type]
                client=client,  # type: ignore[arg-type]
            )
        else:
            request_id, msg, to_send = bwc.batch_command(cmd, ops)
            result = await self.write_command(bwc, cmd, request_id, msg, to_send, client)  # type: ignore[arg-type]

        return result, to_send  # type: ignore[return-value]

    async def _execute_command(
        self,
        generator: Iterator[Any],
        write_concern: WriteConcern,
        session: Optional[AsyncClientSession],
        conn: AsyncConnection,
        op_id: int,
        retryable: bool,
        full_result: MutableMapping[str, Any],
        final_write_concern: Optional[WriteConcern] = None,
    ) -> None:
        db_name = self.collection.database.name
        client = self.collection.database.client
        listeners = client._event_listeners

        if not self.current_run:
            self.current_run = next(generator)
            self.next_run = None
        run = self.current_run

        # AsyncConnection.command validates the session, but we use
        # AsyncConnection.write_command
        conn.validate_session(client, session)
        last_run = False

        while run:
            if not self.retrying:
                self.next_run = next(generator, None)
                if self.next_run is None:
                    last_run = True

            cmd_name = _COMMANDS[run.op_type]
            bwc = self.bulk_ctx_class(
                db_name,
                cmd_name,
                conn,
                op_id,
                listeners,
                session,
                run.op_type,
                self.collection.codec_options,
            )

            while run.idx_offset < len(run.ops):
                # If this is the last possible operation, use the
                # final write concern.
                if last_run and (len(run.ops) - run.idx_offset) == 1:
                    write_concern = final_write_concern or write_concern

                cmd = {cmd_name: self.collection.name, "ordered": self.ordered}
                if self.comment:
                    cmd["comment"] = self.comment
                _csot.apply_write_concern(cmd, write_concern)
                if self.bypass_doc_val:
                    cmd["bypassDocumentValidation"] = True
                if self.let is not None and run.op_type in (_DELETE, _UPDATE):
                    cmd["let"] = self.let
                if session:
                    # Start a new retryable write unless one was already
                    # started for this command.
                    if retryable and not self.started_retryable_write:
                        session._start_retryable_write()
                        self.started_retryable_write = True
                    session._apply_to(cmd, retryable, ReadPreference.PRIMARY, conn)
                conn.send_cluster_time(cmd, session, client)
                conn.add_server_api(cmd)
                # CSOT: apply timeout before encoding the command.
                conn.apply_timeout(client, cmd)
                ops = islice(run.ops, run.idx_offset, None)

                # Run as many ops as possible in one command.
                if write_concern.acknowledged:
                    result, to_send = await self._execute_batch(bwc, cmd, ops, client)

                    # Retryable writeConcernErrors halt the execution of this run.
                    wce = result.get("writeConcernError", {})
                    if wce.get("code", 0) in _RETRYABLE_ERROR_CODES:
                        # Synthesize the full bulk result without modifying the
                        # current one because this write operation may be retried.
                        full = copy.deepcopy(full_result)
                        _merge_command(run, full, run.idx_offset, result)
                        _raise_bulk_write_error(full)

                    _merge_command(run, full_result, run.idx_offset, result)

                    # We're no longer in a retry once a command succeeds.
                    self.retrying = False
                    self.started_retryable_write = False

                    if self.ordered and "writeErrors" in result:
                        break
                else:
                    to_send = await self._execute_batch_unack(bwc, cmd, ops, client)

                run.idx_offset += len(to_send)

            # We're supposed to continue if errors are
            # at the write concern level (e.g. wtimeout)
            if self.ordered and full_result["writeErrors"]:
                break
            # Reset our state
            self.current_run = run = self.next_run

    async def execute_command(
        self,
        generator: Iterator[Any],
        write_concern: WriteConcern,
        session: Optional[AsyncClientSession],
        operation: str,
    ) -> dict[str, Any]:
        """Execute using write commands."""
        # nModified is only reported for write commands, not legacy ops.
        full_result = {
            "writeErrors": [],
            "writeConcernErrors": [],
            "nInserted": 0,
            "nUpserted": 0,
            "nMatched": 0,
            "nModified": 0,
            "nRemoved": 0,
            "upserted": [],
        }
        op_id = _randint()

        async def retryable_bulk(
            session: Optional[AsyncClientSession], conn: AsyncConnection, retryable: bool
        ) -> None:
            await self._execute_command(
                generator,
                write_concern,
                session,
                conn,
                op_id,
                retryable,
                full_result,
            )

        client = self.collection.database.client
        _ = await client._retryable_write(
            self.is_retryable,
            retryable_bulk,
            session,
            operation,
            bulk=self,  # type: ignore[arg-type]
            operation_id=op_id,
        )

        if full_result["writeErrors"] or full_result["writeConcernErrors"]:
            _raise_bulk_write_error(full_result)
        return full_result

    async def execute_op_msg_no_results(
        self, conn: AsyncConnection, generator: Iterator[Any]
    ) -> None:
        """Execute write commands with OP_MSG and w=0 writeConcern, unordered."""
        db_name = self.collection.database.name
        client = self.collection.database.client
        listeners = client._event_listeners
        op_id = _randint()

        if not self.current_run:
            self.current_run = next(generator)
        run = self.current_run

        while run:
            cmd_name = _COMMANDS[run.op_type]
            bwc = self.bulk_ctx_class(
                db_name,
                cmd_name,
                conn,
                op_id,
                listeners,
                None,
                run.op_type,
                self.collection.codec_options,
            )

            while run.idx_offset < len(run.ops):
                cmd = {
                    cmd_name: self.collection.name,
                    "ordered": False,
                    "writeConcern": {"w": 0},
                }
                conn.add_server_api(cmd)
                ops = islice(run.ops, run.idx_offset, None)
                # Run as many ops as possible.
                to_send = await self._execute_batch_unack(bwc, cmd, ops, client)
                run.idx_offset += len(to_send)
            self.current_run = run = next(generator, None)

    async def execute_command_no_results(
        self,
        conn: AsyncConnection,
        generator: Iterator[Any],
        write_concern: WriteConcern,
    ) -> None:
        """Execute write commands with OP_MSG and w=0 WriteConcern, ordered."""
        full_result = {
            "writeErrors": [],
            "writeConcernErrors": [],
            "nInserted": 0,
            "nUpserted": 0,
            "nMatched": 0,
            "nModified": 0,
            "nRemoved": 0,
            "upserted": [],
        }
        # Ordered bulk writes have to be acknowledged so that we stop
        # processing at the first error, even when the application
        # specified unacknowledged writeConcern.
        initial_write_concern = WriteConcern()
        op_id = _randint()
        try:
            await self._execute_command(
                generator,
                initial_write_concern,
                None,
                conn,
                op_id,
                False,
                full_result,
                write_concern,
            )
        except OperationFailure:
            pass

    async def execute_no_results(
        self,
        conn: AsyncConnection,
        generator: Iterator[Any],
        write_concern: WriteConcern,
    ) -> None:
        """Execute all operations, returning no results (w=0)."""
        if self.uses_collation:
            raise ConfigurationError("Collation is unsupported for unacknowledged writes.")
        if self.uses_array_filters:
            raise ConfigurationError("arrayFilters is unsupported for unacknowledged writes.")
        # Guard against unsupported unacknowledged writes.
        unack = write_concern and not write_concern.acknowledged
        if unack and self.uses_hint_delete and conn.max_wire_version < 9:
            raise ConfigurationError(
                "Must be connected to MongoDB 4.4+ to use hint on unacknowledged delete commands."
            )
        if unack and self.uses_hint_update and conn.max_wire_version < 8:
            raise ConfigurationError(
                "Must be connected to MongoDB 4.2+ to use hint on unacknowledged update commands."
            )
        # Cannot have both unacknowledged writes and bypass document validation.
        if self.bypass_doc_val:
            raise OperationFailure(
                "Cannot set bypass_document_validation with unacknowledged write concern"
            )

        if self.ordered:
            return await self.execute_command_no_results(conn, generator, write_concern)
        return await self.execute_op_msg_no_results(conn, generator)

    async def execute(
        self,
        write_concern: WriteConcern,
        session: Optional[AsyncClientSession],
        operation: str,
    ) -> Any:
        """Execute operations."""
        if not self.ops:
            raise InvalidOperation("No operations to execute")
        if self.executed:
            raise InvalidOperation("Bulk operations can only be executed once.")
        self.executed = True
        write_concern = write_concern or self.collection.write_concern
        session = _validate_session_write_concern(session, write_concern)

        if self.ordered:
            generator = self.gen_ordered()
        else:
            generator = self.gen_unordered()

        client = self.collection.database.client
        if not write_concern.acknowledged:
            async with await client._conn_for_writes(session, operation) as connection:
                await self.execute_no_results(connection, generator, write_concern)
                return None
        else:
            return await self.execute_command(generator, write_concern, session, operation)
