"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Tuesday, 21st Jan., 2025.

    OBJECTIVE:

        Live market ticks are broadcasted over Kafka. Here, we try to capture and save those ticks to MongoDb for
        whatever they can be used for later.

    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

# Controllers:
from controllers_v2.finstitutions.trading.all_trading import AllTradingController

# Models:
from models.core.auth_token import CoreAuthTokenModel
from models.finstitutions.trading.symbols import TradingSymbol
from models.finstitutions.trading.ticks import TradingTick
from models.core.user import CoreUserInfoModel

# To work with Zerodha's Kite platform:
from kiteconnect import KiteConnect, KiteTicker

# To work with SocketIO:
import socket
import socketio

# For asynchronous activities:
import asyncio

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

# For MongoDB:
from bson.objectid import ObjectId

# Debugging:
from icecream import IceCreamDebugger


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


# Debugging:
printer = IceCreamDebugger(prefix = "Tick-Save | ", includeContext = True)
no_context_printer = IceCreamDebugger(prefix = "Tick-Save | ", 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())


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

# Redis:
redis_cache: AsyncRedisCache | None = None

# For database:
data_mongo: AsyncMongo | 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)
ticks_since_telegram = 0

# For Zerodha-Kite:
ZERODHA_INSTRUMENT_TOKENS = []
ZERODHA_INSTRUMENT_LOOKUP = {}


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


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

    global ticks_since_telegram

    # Parse the date-time in UTC and add the metadata:
    for tick in ticks:
        tick["metadata"] = {"broker": tick["broker"], "brokerToken": tick["brokerToken"]}
        tick["rcvdTs"] = date_time.parse_date_time(tick["rcvdTs"], timezone = date_time.TIMEZONE_UTC)
        tick["tradeTs"] = date_time.parse_date_time(tick["tradeTs"], timezone = date_time.TIMEZONE_UTC)
        tick["exchgTs"] = date_time.parse_date_time(tick["exchgTs"], timezone = date_time.TIMEZONE_UTC)

    # Save the ticks:
    tasks = [
        data_mongo.insert_many(collection = "__hot_zerodhaTicks", documents = ticks),
        data_mongo.insert_many(collection = "__cold_zerodhaTicks", documents = ticks)
    ]
    results = await asyncio.gather(*tasks)

    # # Save the ticks:
    # inserted_ids = await data_mongo.insert_many(
    #     collection = "__cold_zerodhaTicks",
    #     documents = ticks
    # )

    # Send out the alert if needed:
    ticks_since_telegram += len(ticks)
    ticks_threshold = 1_00_000
    if ticks_since_telegram > ticks_threshold:
        await send_telegram(message = f"*Tick-Save Alert*\nSaved {ticks_since_telegram:,} tick(s).")
        ticks_since_telegram = 0


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


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:
        messages = await consumer.consume(
            count = fetch_count,
            timeout = fetch_timeout
        )

        # If there are no updates to give:
        if not messages:
            no_context_printer("No ticks.")
            continue

        # We extract all the ticks from the Kafka messages:
        ticks = []
        for m in messages:
            if isinstance(m.value, list): ticks += m.value
            else: ticks.append(m.value)

        # Save the ticks:
        no_context_printer(len(ticks))
        await save_ticks(ticks)


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


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 data_mongo
    global redis_cache
    global kafka_consumer

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

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

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

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

    data_mongo = AsyncMongo(
        connection_string = script_cred["mongoDb"]["data"]["connectionString"],
        database_name = "markets",  # script_cred["mongoDb"]["data"]["dbName"],
        max_connections = script_cred["mongoDb"]["data"]["poolSize"],
        debug = debug
    )
    await data_mongo.connect()

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

    redis_cache = AsyncRedisCache(
        connection_string = script_cred["redisCache"]["general"]["connectionString"],
        serializer = JSONSerializer(),
        debug = debug,
        debug_prefix = "General Cache | "
    )
    await redis_cache.connect()

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

    # Create the consumer that will listen to changes in watchlist:
    consumer_creds = script_cred["kafka"]["consumer"]
    kafka_consumer = ConsumerKafka(
        topic = consumer_creds["topic"],
        group_id = "tick_save",
        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.")

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

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


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


async def send_telegram(
        message: str,
        chat_id: str = None,
        message_type: Literal["info", "warning", "error"] = "info"
):

    """
    To send out alerts and heartbeats to inform about th script being alive.
    :param message: The text to send.
    :param chat_id: The destination chat identifier.
    :param message_type: The kind of message to send. Decides the presentation of the header.
    :return:
    """

    try:

        # Create the JSON for sending to the API endpoint:
        json_input = {
            "chatClient": "telegram",
            "message": message,
            "type": message_type
        }
        if chat_id: json_input["chatId"] = json_input

        # Make the API call to send the ticks:
        response = await http_client.post(
            url = r"https://api.thecaoffice.com/converse/tech/alert/chat/backend",
            json = json_input
        )

        # Raise an exception if the call was not successful:
        response.raise_for_status()

    # If something goes wrong:
    except Exception as exception:
        printer(exception)


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


async def heartbeat(interval_seconds: float = 300):

    while True:
        await send_telegram(message = f"*Tick-Save Heartbeat*\nInterval: `{interval_seconds:,} seconds`")
        await asyncio.sleep(interval_seconds)


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


async def main(debug: bool = False):

    tasks = [
        ticks_from_kafka(
            consumer = kafka_consumer,
            fetch_count = 500,
            fetch_timeout = 2.5
        ),
        heartbeat(interval_seconds = 1_800)
    ]

    await asyncio.gather(*tasks)


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


if __name__ == "__main__":

    printer("Main.")

    # To get args from the terminal:
    import argparse

    # 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(
        "-s", "--script-id",
        type = str,
        help = "The id of this script (will affect the loaded config)."
    )
    # parser.add_argument(
    #     "-t", "--token-key", "--token-id",
    #     dest = "token_key",
    #     type = str,
    #     help = "The 'key' to use to retrieve the auth-token for accessing the broker account."
    # )
    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()

    async def runner():

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

        # Initialize and then run the script:
        if await init(
            script_id = args.script_id,
            debug = args.debug
        ): await main(
            debug = args.debug
        )

    asyncio.run(runner())
