# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# NOTE: This code is duplicated from the ingest service:
# src/nv_ingest/util/message_brokers/simple_message_broker/simple_client.py
# Eventually we should move all client wrappers for the message broker into a shared library that both the ingest
# service and the client can use.

import socket
import json
import time
import logging
from typing import Optional, Tuple, Union

from nv_ingest_api.internal.schemas.message_brokers.response_schema import ResponseSchema
from nv_ingest_api.util.service_clients.client_base import MessageBrokerClientBase, FetchMode

logger = logging.getLogger(__name__)


class SimpleClient(MessageBrokerClientBase):
    """
    A client for interfacing with SimpleMessageBroker, creating a new socket connection per request
    to ensure thread safety and robustness. Respects timeouts for all operations.
    """

    def __init__(
        self,
        host: str,
        port: int,
        db: int = 0,
        max_retries: int = 3,
        max_backoff: int = 32,
        connection_timeout: int = 300,
        max_pool_size: int = 128,
        use_ssl: bool = False,
    ):
        """
        Initialize the SimpleClient with configuration parameters.

        Parameters
        ----------
        host : str
            The hostname or IP address of the broker.
        port : int
            The port number of the broker.
        db : int, optional
            The database index (default: 0).
        max_retries : int, optional
            Maximum number of retries for operations (default: 3).
        max_backoff : int, optional
            Maximum backoff time in seconds for retries (default: 32).
        connection_timeout : int, optional
            Timeout in seconds for establishing a connection (default: 300).
        max_pool_size : int, optional
            Maximum pool size for connections (default: 128).
        use_ssl : bool, optional
            Whether to use SSL for connections (default: False).
        """

        self._host = host
        self._port = port
        self._db = db
        self._max_retries = max_retries
        self._max_backoff = max_backoff
        self._max_pool_size = max_pool_size
        self._connection_timeout = connection_timeout
        self._use_ssl = use_ssl

    def get_client(self):
        """
        Retrieve the current client instance.

        Returns
        -------
        SimpleClient
            The current client instance.
        """
        return self

    def submit_message(
        self,
        queue_name: str,
        message: str,
        timeout: Optional[Tuple[int, Union[float]]] = (100, None),
        for_nv_ingest: bool = False,
    ) -> ResponseSchema:
        """
        Submit a message to the specified queue.

        Parameters
        ----------
        queue_name : str
            The name of the queue.
        message : str
            The message to be submitted.
        timeout : float, optional
            Timeout in seconds for the operation.
        for_nv_ingest : bool, optional
            Indicates whether the message is for NV ingest operations.

        Returns
        -------
        ResponseSchema
            The response from the broker.
        """
        return self._handle_push(queue_name, message, timeout, for_nv_ingest)

    def fetch_message(
        self,
        queue_name: str,
        timeout: Optional[Tuple[int, Union[float]]] = (100, None),
        override_fetch_mode: FetchMode = None,
    ) -> ResponseSchema:
        """
        Fetch a message from the specified queue.

        Parameters
        ----------
        queue_name : str
            The name of the queue.
        timeout : float, optional
            Timeout in seconds for the operation.

        Returns
        -------
        ResponseSchema
            The response containing the fetched message.
        """
        if isinstance(timeout, int):
            timeout = (timeout, None)

        return self._handle_pop(queue_name, timeout)

    def ping(self) -> ResponseSchema:
        """
        Ping the broker to check connectivity.

        Returns
        -------
        ResponseSchema
            The response indicating the success of the ping operation.
        """
        command = {"command": "PING"}
        return self._execute_simple_command(command)

    def size(self, queue_name: str) -> ResponseSchema:
        """
        Fetch the size of the specified queue.

        Parameters
        ----------
        queue_name : str
            The name of the queue.

        Returns
        -------
        ResponseSchema
            The response containing the queue size.
        """
        command = {"command": "SIZE", "queue_name": queue_name}
        return self._execute_simple_command(command)

    def _handle_push(
        self, queue_name: str, message: str, timeout: Optional[Tuple[int, Union[float, None]]], for_nv_ingest: bool
    ) -> ResponseSchema:
        """
        Push a message to the queue with optional timeout.

        Parameters
        ----------
        queue_name : str
            The name of the queue.
        message : str
            The message to push.
        timeout : float, optional
            Timeout in seconds for the operation.
        for_nv_ingest : bool
            Indicates whether the message is for NV ingest operations.

        Returns
        -------
        ResponseSchema
            The response from the broker.
        """

        if not queue_name or not isinstance(queue_name, str):
            return ResponseSchema(response_code=1, response_reason="Invalid queue name.")
        if not message or not isinstance(message, str):
            return ResponseSchema(response_code=1, response_reason="Invalid message.")

        if for_nv_ingest:
            command = {"command": "PUSH_FOR_NV_INGEST", "queue_name": queue_name, "message": message}
        else:
            command = {"command": "PUSH", "queue_name": queue_name, "message": message}

        timeout = int(timeout[0])
        if timeout is not None:
            command["timeout"] = timeout

        start_time = time.time()
        while True:
            elapsed = time.time() - start_time
            remaining_timeout = (timeout - elapsed) if (timeout is not None) else None
            if (remaining_timeout is not None) and (remaining_timeout <= 0):
                return ResponseSchema(response_code=1, response_reason="PUSH operation timed out.")

            try:
                with socket.create_connection((self._host, self._port), timeout=self._connection_timeout) as sock:
                    self._send(sock, json.dumps(command).encode("utf-8"))
                    # Receive initial response with transaction ID
                    response_data = self._recv(sock)
                    response = json.loads(response_data)

                    if response.get("response_code") != 0:
                        if (
                            response.get("response_reason") == "Queue is full"
                            or response.get("response_reason") == "Queue is not available"
                        ):
                            time.sleep(0.5)
                            continue
                        else:
                            return ResponseSchema(**response)

                    if "transaction_id" not in response:
                        error_msg = "No transaction_id in response."
                        logger.error(error_msg)

                        return ResponseSchema(response_code=1, response_reason=error_msg)

                    transaction_id = response["transaction_id"]

                    # Send ACK
                    ack_data = json.dumps({"transaction_id": transaction_id, "ack": True}).encode("utf-8")
                    self._send(sock, ack_data)

                    # Receive final response
                    final_response_data = self._recv(sock)
                    final_response = json.loads(final_response_data)

                    return ResponseSchema(**final_response)

            except (ConnectionError, socket.error, BrokenPipeError):
                pass
            except json.JSONDecodeError:
                return ResponseSchema(response_code=1, response_reason="Invalid JSON response from server.")
            except Exception as e:
                return ResponseSchema(response_code=1, response_reason=str(e))

            time.sleep(0.5)  # Backoff delay before retry

    def _handle_pop(self, queue_name: str, timeout: Optional[Tuple[int, Union[float, None]]]) -> ResponseSchema:
        """
        Pop a message from the queue with optional timeout.

        Parameters
        ----------
        queue_name : str
            The name of the queue.
        timeout : float, optional
            Timeout in seconds for the operation.

        Returns
        -------
        ResponseSchema
            The response containing the popped message.
        """

        if not queue_name or not isinstance(queue_name, str):
            return ResponseSchema(response_code=1, response_reason="Invalid queue name.")

        command = {"command": "POP", "queue_name": queue_name}

        timeout = int(timeout[0])

        if timeout is not None:
            command["timeout"] = timeout

        start_time = time.time()
        while True:
            elapsed = time.time() - start_time
            remaining_timeout = timeout - elapsed if timeout else None
            if remaining_timeout is not None and remaining_timeout <= 0:
                return ResponseSchema(response_code=1, response_reason="POP operation timed out.")

            try:
                with socket.create_connection((self._host, self._port), timeout=self._connection_timeout) as sock:
                    self._send(sock, json.dumps(command).encode("utf-8"))
                    # Receive initial response with transaction ID and message
                    response_data = self._recv(sock)
                    response = json.loads(response_data)

                    if response.get("response_code") != 0:
                        if response.get("response_reason") == "Queue is empty":
                            time.sleep(0.1)
                            continue
                        else:
                            return ResponseSchema(**response)

                    if "transaction_id" not in response:
                        error_msg = "No transaction_id in response."

                        return ResponseSchema(response_code=1, response_reason=error_msg)

                    transaction_id = response["transaction_id"]
                    message = response.get("response")

                    # Send ACK
                    ack_data = json.dumps({"transaction_id": transaction_id, "ack": True}).encode("utf-8")
                    self._send(sock, ack_data)

                    # Receive final response
                    final_response_data = self._recv(sock)
                    final_response = json.loads(final_response_data)

                    if final_response.get("response_code") == 0:
                        return ResponseSchema(response_code=0, response=message, transaction_id=transaction_id)
                    else:
                        return ResponseSchema(**final_response)

            except (ConnectionError, socket.error, BrokenPipeError):
                pass
            except json.JSONDecodeError:
                return ResponseSchema(response_code=1, response_reason="Invalid JSON response from server.")
            except Exception as e:
                return ResponseSchema(response_code=1, response_reason=str(e))

            time.sleep(0.1)  # Backoff delay before retry

    def _execute_simple_command(self, command: dict) -> ResponseSchema:
        """
        Execute a simple command on the broker and process the response.

        Parameters
        ----------
        command : dict
            The command to execute.

        Returns
        -------
        ResponseSchema
            The response from the broker.
        """

        if isinstance(command, dict):
            data = json.dumps(command).encode("utf-8")
        elif isinstance(command, str):
            data = command.encode("utf-8")

        try:
            with socket.create_connection((self._host, self._port), timeout=self._connection_timeout) as sock:
                self._send(sock, data)
                response_data = self._recv(sock)
                response = json.loads(response_data)
                return ResponseSchema(**response)
        except (ConnectionError, socket.error, BrokenPipeError) as e:
            return ResponseSchema(response_code=1, response_reason=f"Connection error: {e}")
        except json.JSONDecodeError:
            return ResponseSchema(response_code=1, response_reason="Invalid JSON response from server.")
        except Exception as e:
            return ResponseSchema(response_code=1, response_reason=str(e))

    def _send(self, sock: socket.socket, data: bytes) -> None:
        """
        Send data over a socket connection with a length header.

        Parameters
        ----------
        sock : socket.socket
            The socket connection.
        data : bytes
            The data to send.

        Raises
        ------
        ConnectionError
            If sending data fails.
        """

        total_length = len(data)
        if total_length == 0:
            raise ValueError("Cannot send an empty message.")

        try:
            sock.sendall(total_length.to_bytes(8, "big"))
            sock.sendall(data)
        except (socket.error, BrokenPipeError):
            raise ConnectionError("Failed to send data.")

    def _recv(self, sock: socket.socket) -> str:
        """
        Receive data from a socket connection based on a length header.

        Parameters
        ----------
        sock : socket.socket
            The socket connection.

        Returns
        -------
        str
            The received data as a string.

        Raises
        ------
        ConnectionError
            If receiving data fails.
        """

        try:
            length_header = self._recv_exact(sock, 8)
            if not length_header:
                raise ConnectionError("Incomplete length header received.")
            total_length = int.from_bytes(length_header, "big")
            data_bytes = self._recv_exact(sock, total_length)
            if not data_bytes:
                raise ConnectionError("Incomplete message received.")
            return data_bytes.decode("utf-8")
        except (socket.error, BrokenPipeError, ConnectionError):
            raise ConnectionError("Failed to receive data.")

    def _recv_exact(self, sock: socket.socket, num_bytes: int) -> Optional[bytes]:
        """
        Receive an exact number of bytes from a socket connection.

        Parameters
        ----------
        sock : socket.socket
            The socket connection.
        num_bytes : int
            The number of bytes to receive.

        Returns
        -------
        Optional[bytes]
            The received bytes, or None if the connection is closed.
        """

        data = bytearray()
        while len(data) < num_bytes:
            try:
                packet = sock.recv(num_bytes - len(data))
                if not packet:
                    return None
                data.extend(packet)
            except socket.timeout:
                return None
            except Exception:
                return None
        return bytes(data)
