"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 12th Dec., 2024

    OBJECTIVE:

        To handle all messages from one place.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""

#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# My async utils:
from utils_v2.string import json
from utils_v2.date_time import date_time
from utils_v2.database.async_mongo_v2 import AsyncMongo, AsyncMongoStorage

# Models:
from models.servers.core import CoreServerInfoModel, CoreServerCheckLogModel

# To make API calls:
import httpx

# To work with MongoDB:
from bson import ObjectId
from pymongo import InsertOne, UpdateOne, ReplaceOne

# To work with datatypes:
from typing import Literal, List, Dict, Any

# To make deep-copies:
import copy

# To work with base-64 encoding:
import base64

# To work with date and time:
import datetime
import time

# For asynchronous activities:
import asyncio

# For debugging:
from icecream import IceCreamDebugger


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class CoreServerController:

    # ┏┓┓      ┓┏
    # ┃ ┃┏┓┏┏  ┃┃┏┓┏┓┏
    # ┗┛┗┗┻┛┛  ┗┛┗┻┛ ┛

    # For MongoDB:
    SERVERS_COLLECTION = "_servers"
    SERVER_CHECK_LOGS_COLLECTION = "_serverCheckLogs"

    # ┏┓
    # ┃ ┏┓┏┓┏╋┏┓┓┏┏╋┏┓┏┓
    # ┗┛┗┛┛┗┛┗┛ ┗┻┗┗┗┛┛

    def __init__(
            self,
            batch_size: int = 50,
            batch_timeout: int | float = 60,
            http_client: httpx.AsyncClient = None,
            debug = True,
            debug_prefix = "Srvr. (C) | ",
            debug_only_errors = True
    ):

        """
        To initialize the instance of this Server controller.
        :param batch_size: When monitoring the servers, how many records will one worker pick at once to poll.
        :param batch_timeout: When a particular server has been picked in a batch, for how long must another worker not
            touch it before making his own attempt.
        :param http_client: An instance of 'httpx' library's AsyncClient. If not given, one will be instantiated
            internally. It is recommended that, for multi-bot use cases, you provide a common HTTP client from outside.
        :param debug: Whether, or not, you would like to show debugging messages on the terminal.
        :param debug_prefix: The prefix string to identify the debugging messages.
        :param debug_only_errors: Whether you would like to show all debugging messages or just error messages.
        """

        # Prepare the debugging utility:
        self._debug_prefix = debug_prefix
        self._printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: self._printer.disable()
        self._debug_only_errors = debug_only_errors

        # Accept the configuration:
        self._batch_size = batch_size
        self._batch_timeout = batch_timeout

        # Accept/create an HTTP client to work with:
        if http_client: self.__http_client = http_client
        else: self.__http_client = httpx.AsyncClient(
            limits = httpx.Limits(
                max_connections = 100,  # ............ Maximum number of connections allowed in the pool.
                max_keepalive_connections = 50,  # ... Maximum number of connections that can be kept alive.
            ),
            timeout = httpx.Timeout(
                pool = 120.0,  # .... Time to wait for a free connection from the pool.
                connect = 2.5,  # ... Time to wait for establishing a connection to the server.
                write = 5.0,  # ..... Time to wait for sending data.
                read = 5.0  # ....... Time to wait for receiving data.
            )
        )

    # ┏┓┳┓┳┳┳┓      ┏┓
    # ┃ ┣┫┃┃┃┃  ━━  ┃ ┏┓┏┓┏┓╋┏┓
    # ┗┛┛┗┗┛┻┛      ┗┛┛ ┗ ┗┻┗┗

    async def register(
            self,
            mongo_data_conn: AsyncMongo,
            server: CoreServerInfoModel
    ) -> CoreServerInfoModel | None:

        """
        Register one server in the database.
        :param mongo_data_conn: The instance of the database connector to use for the operation.
        :param server: The message to save into the database.
        :return: The object id of the inserted document.
        """

        # Note down the current time:
        event_ts = date_time.get_current_utc_date_time(as_string = False)

        # Register/update the server's info:
        _server = await mongo_data_conn.find_one_and_update(
            collection = self.SERVERS_COLLECTION,
            filter = {
                "hostname": server.hostname,
                "project": server.project,
                "service": server.service,
                "healthCheckUrl": server.healthCheckUrl
            },
            update = {
                "$setOnInsert": {
                    "hostname": server.hostname,
                    "project": server.project,
                    "service": server.service,
                    "healthCheckUrl": server.healthCheckUrl,
                    "firstRegTs": event_ts
                },
                "$set": {
                    "os": server.os,
                    "cpu": server.cpu,
                    "pid": server.pid,
                    "ppid": server.ppid,
                    "ipAddr": server.ipAddr,
                    "portNo": server.portNo,
                    "description": server.description,
                    "healthCheckInterval": server.healthCheckInterval,
                    "healthAlertUrl": server.healthAlertUrl,
                    "online": True,
                    "batchId": server.batchId,
                    "batchTs": server.batchTs,
                    "lastRegTs": server.lastRegTs,
                    "lastCheckTs": server.lastCheckTs,
                    "checkAfterTs": server.checkAfterTs,
                }
            },
            upsert = True,
            return_updated = True
        )

        # Done here:
        return CoreServerInfoModel(**_server) if _server else None

    # ┏┓┳┓┳┳┳┓      ┳┓     •
    # ┃ ┣┫┃┃┃┃  ━━  ┣┫┏┓╋┏┓┓┏┓┓┏┏┓
    # ┗┛┛┗┗┛┻┛      ┛┗┗ ┗┛ ┗┗ ┗┛┗

    async def get_batch(
            self,
            mongo_data_conn: AsyncMongo,
    ) -> List[CoreServerInfoModel] | None:

        # Get a batch id and note down the time:
        batch_id = ObjectId()
        batch_ts = date_time.get_current_utc_date_time(as_string = False)

        # Find servers that need to be checked:
        servers = await mongo_data_conn.find_many(
            collection = self.SERVERS_COLLECTION,
            filter = {
                "checkAfterTs": {"$lte": batch_ts},
                "$or": [
                    {
                        "batchId": {"$eq": None}
                    },
                    {
                        "batchTs": {"$lt": batch_ts - datetime.timedelta(seconds = self._batch_timeout)}
                    }
                ]
            },
            sort = {"checkAfterTs": 1},
            limit = self._batch_size
        )

        # If we didn't find any servers:
        if not servers: return servers

        # Model the servers:
        servers = [CoreServerInfoModel(**s) for s in servers]

        # Mark these servers as picked for checking:
        updated_count = await mongo_data_conn.update_many(
            collection = self.SERVERS_COLLECTION,
            filter = {"_id": {"$in": [s.serverId for s in servers]}},
            update = {
                "$set": {
                    "batchId": batch_id,
                    "batchTs": batch_ts
                }
            },
            upsert = False
        )

        # Done here:
        if updated_count < len(servers):
            self._printer("Mismatched count", len(servers), updated_count)
            return None
        else: return servers

    # ┏┓┳┓┳┳┳┓      ┳┳   ┓
    # ┃ ┣┫┃┃┃┃  ━━  ┃┃┏┓┏┫┏┓╋┏┓
    # ┗┛┛┗┗┛┻┛      ┗┛┣┛┗┻┗┻┗┗
    #                 ┛

    async def check_one_server(
            self,
            mongo_data_conn: AsyncMongo,
            server: CoreServerInfoModel,
            insert_check_log: bool = True,
            release_from_batch: bool = True
    ) -> CoreServerCheckLogModel:

        # Note down some assumptions and starting parameters:
        server_check_log = CoreServerCheckLogModel(serverId = server.serverId)
        start_time = time.time()

        try:

            # Make the API call:
            server_response = await self.__http_client.get(url = server.healthCheckUrl)
            server_response.raise_for_status()
            server_check_log.online = True
            server_check_log.message = "The server responded."

        # Catch various exceptions:
        except httpx.HTTPStatusError as e: server_check_log.message = f"HTTP Error: {e.response.status_code}"
        except httpx.ConnectTimeout as e: server_check_log.message = f"Connect Timeout Error: {e}"
        except httpx.ReadTimeout as e: server_check_log.message = f"Read Timeout Error: {e}"
        except httpx.TimeoutException as e: server_check_log.message = f"Misc. Timeout Error: {e}"
        except httpx.NetworkError as e: server_check_log.message = f"Network Error: {e}"
        except httpx.RequestError as e: server_check_log.message = f"Request Error: {e}"

        # Note down the latency:
        server_check_log.latency = time.time() - start_time

        # In case of failure, we try to hit the alert URL:
        if not server_check_log.online:
            try:
                alert_url_response = await self.__http_client.get(url = server.healthAlertUrl)
                server_check_log.alertRaised = True if alert_url_response.status_code == 200 else False
            except:
                server_check_log.alertRaised = False

        # Store the log in the database:
        if insert_check_log: await mongo_data_conn.insert_one(
            collection = self.SERVER_CHECK_LOGS_COLLECTION,
            document = server_check_log.model_dump()
        )

        # Release the picked server from the batch:
        if release_from_batch: await mongo_data_conn.update_one(
            collection = self.SERVERS_COLLECTION,
            filter = {"_id": server.serverId},
            update = {
                "$set": {
                    "batchId": None,
                    "batchTs": None,
                    "online": server_check_log.online,
                    "lastCheckTs": server_check_log.ts,
                    "checkAfterTs": server_check_log.ts + datetime.timedelta(seconds = server.healthCheckInterval)
                }
            },
            upsert = False
        )
        
        # Done here:
        return server_check_log

    async def check_server_batch(
            self,
            mongo_data_conn: AsyncMongo,
            servers: List[CoreServerInfoModel]
    ) -> List[CoreServerCheckLogModel]:

        # If the servers list is a blank array:
        if not servers: return []

        # Create tasks to check individual servers,
        # tell the tasks not to update the database individually:
        tasks = [
            self.check_one_server(
                mongo_data_conn = mongo_data_conn,
                server = server,
                insert_check_log = False,
                release_from_batch = False
            ) for server in servers
        ]

        # Fire the tasks and note down the results:
        server_check_logs = await asyncio.gather(*tasks)

        # Insert the logs:
        await mongo_data_conn.insert_many(
            collection = self.SERVER_CHECK_LOGS_COLLECTION,
            documents = [log.model_dump() for log in server_check_logs]
        )

        # Release the batches:
        await mongo_data_conn.bulk_write(
            collection = self.SERVERS_COLLECTION,
            requests = [
                UpdateOne(
                    filter = {"_id": server.serverId},
                    update = {
                        "$set": {
                            "batchId": None,
                            "batchTs": None,
                            "online": log.online,
                            "lastCheckTs": log.ts,
                            "checkAfterTs": log.ts + datetime.timedelta(seconds = server.healthCheckInterval)
                        }
                    },
                    upsert = False
                ) for server, log in zip(servers, server_check_logs)
            ]
        )

        # Done here:
        return server_check_logs


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    pass