"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 27th Dec., 2024

    OBJECTIVE:

        To broadcast live tick updates to connected clients.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


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


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

# System-level activities:
import io
import os

# My utils:
from utils_v2.string import json
from utils_v2.system import files
from utils_v2.database.async_mongo_v2 import AsyncMongo
from utils_v2.queue.async_kafka import ConsumerKafka, get_ssl_context

# To make HTTP calls:
import httpx

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

# Models:
from models.finstitutions.trading.symbols import TradingSymbol
from models.finstitutions.trading.ticks import TradingTick

# To work with SocketIO:
import socketio

# For asynchronous activities:
import asyncio

# To work with various datatypes:
from typing import List

# Debugging:
from icecream import IceCreamDebugger


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


# Debugging:
printer = IceCreamDebugger(prefix = "Tick-Disp, | ", includeContext = True)
printer.disable()

# For SocketIO:
sio = socketio.AsyncServer(async_mode = "asgi")
app = socketio.ASGIApp(sio)

# SocketIO Namespaces:
NAMESPACE_MODULE = "/finstitutions/trading"
NAMESPACE_PASSTHROUGH = "/passthrough"

# SocketIO Events:
EVENT_CONNECT = "connect"
EVENT_DISCONNECT = "disconnect"
EVENT_ECHO = "echo"
EVENT_TICKS = "ticks"


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


# Session-awareness:
pass

# Script-local:
flags = {
    "initDone": False
}


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


@sio.on(event = EVENT_CONNECT, namespace = NAMESPACE_MODULE)
async def handle_connect(sid, environ):

    # Start the common background processes:
    if not flags.get("initDone"):
        flags["initDone"] = True
        asyncio.create_task(init())

    # Allow/reject requests:
    printer(sid)
    print("ENVIRON:", json.to_string(environ, default = str))
    return True


# ---------------------------------------------------------------------------------------------------------------------


@sio.on(event = EVENT_DISCONNECT, namespace = NAMESPACE_MODULE)
async def handle_disconnect(sid, reason):
    printer(sid, reason)


# ---------------------------------------------------------------------------------------------------------------------


@sio.on(event = EVENT_ECHO, namespace = NAMESPACE_MODULE)
async def handle_echo(sid, data):

    """
    For testing. This is a quick way to check if the module is up.
    :param sid: The id of the client that caused this event.
    :param data: The data sent by the client.
    :return: None.
    """

    printer(sid)
    await sio.emit(
        event = EVENT_ECHO,
        data = data,
        namespace = NAMESPACE_MODULE
    )


# ---------------------------------------------------------------------------------------------------------------------


async def send_ticks(ticks: List[dict]):

    """
    Here's where we decide which client gets which tick and send it out.
    WARNING: WE ARE ASSUMING THAT NO FURTHER FORMATING/COMPUTATION IS REQUIRED OTHER THAN SELECTING WHICH SUBSETS OF
             TICKS TO SEND TO WHICH CLIENTS. FOR US THE TICKS ALREADY HAVE ALL THE DATA NEEDED TO BE SEND TO
             RESPECTIVE CLIENTS.
    :param ticks: The list of individual tick updates to send out to the clients.
    :return: ??
    """

    # Currently we're just broadcasting
    # all the data to all the clients:
    await sio.emit(
        event = EVENT_TICKS,
        data = ticks,
        namespace = NAMESPACE_MODULE
    )


# ---------------------------------------------------------------------------------------------------------------------


async def ticks_from_kafka(
        consumer: ConsumerKafka,
        fetch_count: int = 100,
        fetch_timeout: float = 1.0
) -> None:

    """
    This function must run in the background forever and just keep listening for ticks on Kafka and keep relaying them
    to all the connected clients as per their watchlists.
    :param consumer: The preconfigured Kafka consumer that can listen for ticks in asynchronous mode.
    :param fetch_count: How many messages to consume in one go.
    :param fetch_timeout: How long to wait (in seconds) while consuming messages from Kafka.
    :return: None
    """

    # Do the next part infinitely:
    while True:

        # Get messages form Kafka:
        ticks = await consumer.consume(
            count = fetch_count,
            timeout = fetch_timeout
        )

        # If there are no updates to give:
        if not ticks: continue

        # Each message must be treated as an array of tick updates (list of dicts).
        # In case the producer is sending each individual tick as a separate message,
        # we normalize it to be a list:
        tasks = [send_ticks(t["value"] if isinstance(t["value"], list) else [t["value"]]) for t in ticks]
        results = await asyncio.gather(*tasks)
        printer(len(ticks))


# ---------------------------------------------------------------------------------------------------------------------


async def init():

    # Handle debugging:
    if os.environ["DEBUG"].lower() == "true": printer.enable()
    printer("Initializing.")

    # Start consuming ticks in the background:
    cwd = files.get_cwd()
    parent_dir = cwd
    sio.start_background_task(
        ticks_from_kafka,
        consumer = ConsumerKafka(
            topic = "tickers",
            bootstrap_servers = "del.ditscentre.in:9092",
            security_protocol = "SSL",
            ssl_context = get_ssl_context(
                # ca_file = "../../../creds/kafka/cert_authority.pem",
                # cert_file = "../../../creds/kafka/fullchain.pem",
                # key_file = "../../../creds/kafka/privkey.pem"
                ca_file = os.path.join(parent_dir, "creds", "kafka", "cert_authority.pem"),
                cert_file = os.path.join(parent_dir, "creds", "kafka", "fullchain.pem"),
                key_file = os.path.join(parent_dir, "creds", "kafka", "privkey.pem")
            ),
            auto_offset_reset = "latest"
        ),
        fetch_count = 1_250,
        fetch_timeout = 1.0
    )


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


if __name__ == "__main__":

    # To get args from the terminal:
    import argparse

    # To run the ASGI:
    import uvicorn
    from multiprocessing import freeze_support

    # Get the config from the command-line:
    parser = argparse.ArgumentParser(description = f"SocketIO to serve live market data.")
    parser.add_argument(
        "--workers",
        type = int,
        help = "The no. of threads to spin up for this instance!",
        default = 2
    )
    parser.add_argument(
        "--host",
        type = str,
        help = "The host for the app. e.g.: '0.0.0.0' or '127.0.0.1'.",
        default = "127.0.0.1"
    )
    parser.add_argument(
        "--port",
        type = int,
        help = "The port no. to bind the app to.",
        default = 8080
    )
    parser.add_argument(
        "--script-id",
        type = str,
        help = "The id of this script (will affect the loaded config)."
    )
    parser.add_argument(
        "--debug",
        action = "store_true",
        help = "Whether, or not, you want to see debugging messages in the terminal.",
        default = False
    )
    args = parser.parse_args()

    # Note down the config;
    os.environ["SCRIPT_ID"] = args.script_id
    os.environ["DEBUG"] = str(args.debug)

    # Run the gateway:
    freeze_support()
    uvicorn.run(
        app = "main:app",
        workers = args.workers,
        host = args.host,
        port = args.port
    )
