"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 2nd Jan., 2025

    OBJECTIVE:

        To broadcast live tick updates to connected clients. It doesn't matter which stockbroker we are getting the
        ticks from as long as we are reading standardized ticks from the Kafka queue.

    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.string import regex
from utils_v2.system import files
from utils_v2.date_time import date_time
from utils_v2.database.async_mongo_v2 import AsyncMongo
from utils_v2.queue.kafka.controllers.async_kafka import ConsumerKafka, get_ssl_context
from utils_v2.cache.async_redis_cache_v2 import AsyncRedisCache
from utils_v2.serialization.json_serializer import JSONSerializer

# To make HTTP calls:
import httpx

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

# Models:
from models.core.user import CoreUserInfoModel

# To work with SocketIO:
import socket
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-Out | ", includeContext = True)
no_context_printer = IceCreamDebugger(prefix = "Tick-Out | ", includeContext = False)

# To make API calls:
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 = 10.0,  # .... Time to wait for sending data.
        read = 9.9  # ....... Time to wait for receiving data.
    )
)

# General:
SERVER_HOSTNAME = str(socket.gethostname())

# For SocketIO:
# Namespaces:
NAMESPACE_MODULE = None
NAMESPACE_PASSTHROUGH = "/passthrough"
# Events:
EVENT_CONNECT = "connect"
EVENT_DISCONNECT = "disconnect"
EVENT_ECHO = "echo"
EVENT_TICKS = "ticks"


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


# For SocketIO:
ALLOWED_ORIGINS = []
sio = socketio.AsyncServer(
    cors_allowed_origins = "*",
    async_mode = "asgi"
)
app = socketio.ASGIApp(sio)

# Redis:
redis_cache: AsyncRedisCache | None = None

# For kafka:
kafka_consumer: ConsumerKafka | None = None

# Session-awareness and maintenance of this script's state:
SCRIPT_DATA = {}
exclusive_lock = asyncio.Semaphore(1)
CONNECTED_CLIENTS = {}
FLAGS = {
    "initDone": False
}


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


def origin_is_allowed(origin: str) -> bool:

    """
    To check if a given origin is in the allowed list.
    :param origin: The origin of your request.
    :return: True if allowed, else False.
    """

    # Start by assuming failure:
    is_allowed = False

    # Check through all the allowed origins:
    for allowed in ALLOWED_ORIGINS:
        try:
            if regex.match(origin, allowed):
                is_allowed = True
                break
        except Exception as exception:
            printer(exception)

    # Done here:
    return is_allowed


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


async def send_ticks(ticks: List[dict]) -> None:

    """
    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: None
    """

    # 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
    """

    printer("Starting Kafka consumer (ticks).")

    # Do the next part infinitely:
    while True:

        # Note the time:
        now_utc = date_time.get_current_utc_date_time().timestamp()

        # 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)

        # Analyze the ticks:
        # latency = [abs(now_utc - t.value.get("rcvdTs", t.value["tradeTs"])) for t in ticks]
        latency = [abs(now_utc - t.ts.timestamp()) for t in ticks]
        avg_latency = sum(latency) / len(latency)
        total_ticks = len(ticks)
        # late_cutoff_seconds = 3.0
        #
        # late_ticks = 0
        # for tick in ticks:
        #     if now_utc - tick.value["tradeTs"] > late_cutoff_seconds:
        #         late_ticks += 1
        # ticks_str = f"COUNT: {total_ticks: >5,}   |   LATE: {late_ticks: >5,} ({(late_ticks/total_ticks)*100.0:.2f}%)"
        ticks_str = f"COUNT: {total_ticks: >5,}   |   AVG. LATENCY: {avg_latency:.5f}"
        no_context_printer(ticks_str)


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


async def init(
        script_id: str,
        debug: bool
):

    """
    To initialize all credentials, instances, and connectivity for this whole script.
    :param script_id: The id to use to load cred and data from the internal service.
    :param debug: Whether, or not, you would like to print the debug messages.
    :return: True if initialized successfully, else False.
    """

    # Declare the required global variables:
    global SCRIPT_DATA
    global ALLOWED_ORIGINS
    global redis_cache
    global kafka_consumer

    # Basic stuff:
    if debug: printer.enable()
    printer("Initializing.")

    # ┏┓  •  •
    # ┃┃┏┓┓┏┓┓┏┓┏
    # ┗┛┛ ┗┗┫┗┛┗┛
    #       ┛

    response = await http_client.post(
        url = r"https://api.thecaoffice.com/ca/get/title",
        headers = {"Origin": "https://thecaoffice.com/"},
        data = {
            "domainName": "127.0.0.1:1234",
            "screenWidth": 1920,
            "screenHeight": 1080
        }
    )
    if response.status_code not in [200]:
        print("FATAL: ALLOWED ORIGINS NOT FETCHED!")
        return False
    ALLOWED_ORIGINS = [origin["domain"] for origin in response.json().get("data", {}).get("rs2", [])]
    printer(ALLOWED_ORIGINS)
    if len(ALLOWED_ORIGINS) < 1:
        print("FATAL: ALLOWED ORIGINS IS EMPTY!")
        return False

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

    # Get the script credentials:
    response = await http_client.get(
        url = r"https://nexcom.ditscentre.in/internal/cred/get",
        headers = {"X-Script-Id": script_id}
    )
    if response.status_code not in [200]:
        print("FATAL: SCRIPT CREDENTIALS LOADING FAILED!")
        return False
    script_cred = response.json().get("data")

    # Get the script data:
    response = await http_client.get(
        url = r"https://nexcom.ditscentre.in/internal/data/get",
        headers = {"X-Script-Id": script_id}
    )
    if response.status_code not in [200]:
        print("FATAL: SCRIPT DATA LOADING FAILED!")
        return False
    SCRIPT_DATA = response.json().get("data")

    # Done with this step:
    printer("Cred and Data loaded.")

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

    # Create the consumer that will listen to changes in watchlist:
    consumer_creds = script_cred["kafka"]["consumer"]
    kafka_consumer = ConsumerKafka(
        topic = consumer_creds["topic"],
        bootstrap_servers = consumer_creds["config"]["bootstrapServers"],
        security_protocol = consumer_creds["config"].get("securityProtocol", "PLAINTEXT"),
        ssl_context = get_ssl_context(
            ca_file = consumer_creds["config"].get("caFile"),
            cert_file = consumer_creds["config"].get("certFile"),
            key_file = consumer_creds["config"].get("keyFile"),
        ),
        serializer = JSONSerializer(),
        debug = debug
    )
    if not await kafka_consumer.connect():
        print("FATAL: KAFKA CONSUMER NOT CREATED!")
        return False
    printer("Kafka consumer ready.")

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

    redis_cache = AsyncRedisCache(
        connection_string = script_cred["redisCache"]["general"]["connectionString"],
        serializer = JSONSerializer(),
        debug = debug,
        debug_prefix = "General Cache | "
    )
    if not await redis_cache.connect():
        print("FATAL: REDIS CACHE NOT CREATED!")
        return False
    printer("Redis cache ready.")

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

    # Start the background task that will receive ticks from the Kafka queue and broadcast them to the respective
    # connected clients:
    sio.start_background_task(
        ticks_from_kafka,
        consumer = kafka_consumer,
        fetch_count = 1_000,
        fetch_timeout = 1.0
    )

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

    # If everything went well, we return with success:
    printer("Initialization done.")
    return True


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


@sio.on(event = EVENT_CONNECT, namespace = NAMESPACE_MODULE)
async def on_connect(sid, environ, *args) -> bool:

    """
    The event handler for when a new connection request comes in.
    :param sid: The session id of the incoming request (generated by SocketIO).
    :param environ: The set of headers and other connection-specific values.
    :param args: Any extra input coming from the connection request.
    :return: True to accept a connection request, False to reject it.
    """

    # declare the required global variables:
    global CONNECTED_CLIENTS

    # Initialize the script if needed:
    async with exclusive_lock:
        if not FLAGS.get("initDone"):
            FLAGS["initDone"] = await init(
                script_id = os.environ["SCRIPT_ID"],
                debug = True if os.environ["DEBUG"].lower() == "true" else False
            )

    # If the initialization failed, we cannot accept the incoming request:
    if not FLAGS.get("initDone"):
        printer("SOCKET REJECTED: Init. pending.", sid)
        return False

    # Check the origin of the incoming request:
    printer("Checking origin.")
    origin = environ.get("HTTP_ORIGIN", "???")
    if not origin_is_allowed(origin):
        printer("SOCKET REJECTED: Bad origin.", sid, origin)
        return False

    # Get the session token from the incoming request:
    printer("Checking session token.")
    session_token = environ.get("HTTP_X_SESSION_TOKEN")
    if not session_token and len(args) > 0: session_token = args[0].get("X-Session-Token")
    if not session_token:
        printer("SOCKET REJECTED: No session token.", sid)
        return False

    # Get the user's details from the session token:
    printer("Fetching user info.")
    user_info = await redis_cache.get(key = session_token)
    if not user_info:
        printer("SOCKET REJECTED: Invalid session token.", sid)
        return False
    user_info = CoreUserInfoModel(**user_info)

    # Get the user's watchlist and note down the details.
    # Consider the following structure for a user's info:
    redis_key = f"io_{session_token}"
    async with exclusive_lock:
        CONNECTED_CLIENTS[sid] = {
            "user": user_info,
            "redisKey": redis_key,
            "rooms": []
        }
    sid_cached = await redis_cache.set(
        key = redis_key,
        value = {"server": SERVER_HOSTNAME, "socket_id": sid}
    )

    # Done here:
    printer("SOCKET ACCEPTED.", sid, sid_cached)
    return True


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


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

    """
    To handle a disconnect event. Automatically triggered when a client disconnects from the server.
    :param sid: The session id of the client (generated by SocketIO on connecting).
    :param reason: The hint about why the disconnection happened.
    :return: None.
    """

    # declare the required global variables:
    global CONNECTED_CLIENTS

    # register the disconnect in the global variable, and on the cache server:
    client_info = {}
    async with exclusive_lock: client_info = CONNECTED_CLIENTS.pop(sid, None)
    sid_uncached = await redis_cache.delete(key = client_info["redisKey"]) if client_info else False
    printer("SOCKET DISCONNECTED", sid, reason, sid_uncached)


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


if __name__ == "__main__":

    printer("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 (and a general passthrough).")
    parser.add_argument(
        "-w", "--workers",
        type = int,
        help = "The no. of threads to spin up for this instance!",
        default = 2
    )
    parser.add_argument(
        "-a", "--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(
        "-p", "--port",
        type = int,
        help = "The port no. to bind the app to.",
        default = 8080
    )
    parser.add_argument(
        "-s", "--script-id",
        type = str,
        help = "The id of this script (will affect the loaded config)."
    )
    parser.add_argument(
        "-d", "--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)

    # Startup message:
    printer.enable()
    printer(str(args.debug))
    printer.disable()

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

