"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Wednesday, 1st Jan., 2025.

    OBJECTIVE:

        To get live updates from Zerodha and push them to Kafka.

    REFERENCES:

        01. YouTube Webinar: https://www.youtube.com/watch?v=9vzd289Eedk
        02. Official Example (GitHub): https://github.com/zerodha/pykiteconnect/blob/master/examples/threaded_ticker.py

    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.cache.async_redis_cache_v2 import AsyncRedisCache
from utils_v2.queue.kafka.controllers.kafka import ProducerKafka, ConsumerKafka
from utils_v2.serialization.json_serializer import JSONSerializer

# To make HTTP calls:
import httpx
import socket

# 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

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

# To work with datatypes:
from typing import List

# For debugging:
from icecream import IceCreamDebugger

# For asynchronous operations:
import asyncio


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


# To make API calls:
http_client = httpx.Client(
    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.
    )
)

# For debugging:
printer = IceCreamDebugger(prefix = "Tick-In (Sful) | ", includeContext = True)
no_context_printer = IceCreamDebugger(prefix = "Tick-In (Sful) | ", includeContext = False)


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


# To identify this process:
SERVER_HOSTNAME = str(socket.gethostname())
TOKEN_KEY = None
AUTH_TOKEN = None

# Pertaining to the behaviour of this script:
SCRIPT_DATA = {}

# For kafka:
kafka_producer: ProducerKafka | None = None
kafka_consumer: ConsumerKafka | None = None
TOTAL_TICK_COUNT = 0
TICKS_SINCE_FLUSH = 0

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


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


def ticks_to_kafka(ticks: List[TradingTick]) -> int:

    """
    This function simply takes the ticks received from the stockbroker's servers and sends them to the internal Kafka
    queue for various end consumption use cases.
    :param ticks: The modeled ticks that must be pushed to the Kafka queue.
    :return: The count of the ticks that were successfully pushed to the queue.
    """

    # Declare the required global variables:
    global TOTAL_TICK_COUNT
    global TICKS_SINCE_FLUSH

    # Start with basic variables:
    total_count = len(ticks)
    success_count = 0
    failure_count = 0

    # Push out the ticks to the queue:
    for tick in ticks:

        # Flush the existing messages (if needed):
        TOTAL_TICK_COUNT += 1
        TICKS_SINCE_FLUSH += 1
        if TICKS_SINCE_FLUSH > 50_000:
            kafka_producer.flush(timeout = 0.0)
            TICKS_SINCE_FLUSH = 0

        # Push this one tick to the queue:
        success = kafka_producer.produce(value = tick.summary)
        kafka_producer.client.poll(0)
        if success: success_count += 1
        else: failure_count += 1

    # Debugging print:
    ticks_str = f"|   TICKS: {total_count:6,}   |   PRDC'D: {success_count:6,}   |   TOT: {TOTAL_TICK_COUNT:10,}   |"
    no_context_printer(ticks_str)

    # Done here:
    return success_count


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


def on_zerodha_kite_connect(ws, response) -> None:

    """
    To do something when we get connected to Zerodha's Kite API successfully.
    :param ws: The websocket object that is connected to Zerodha.
    :param response: Some input from Zerodha's library.
    :return: None
    """

    printer("Zerodha Kite WS connected.")
    ws.subscribe(ZERODHA_INSTRUMENT_TOKENS)
    ws.set_mode(ws.MODE_FULL, ZERODHA_INSTRUMENT_TOKENS)
    printer("Zerodha Kite instruments subscribed.", len(ZERODHA_INSTRUMENT_TOKENS))


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


def on_zerodha_kite_ticks(ws, ticks) -> None:

    """
    The function that gets called when Zerodha's tick updates come in.
    :param ws: The websocket object that is connected to Zerodha.
    :param ticks: The actual tick data received from Zerodha's Kite platform.
    :return: None
    """

    # Model the raw input ticks:
    ticks = TradingTick.from_zerodha_kite(
        ticks = ticks,
        instrument_lookup = ZERODHA_INSTRUMENT_LOOKUP,
        received_ts = date_time.get_current_utc_date_time(as_string = False)
    )

    # Send the ticks to Kafka:
    success_count = ticks_to_kafka(ticks)


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


def setup_zerodha_kite_feed(
        auth_token: CoreAuthTokenModel,
        total_instruments: List[dict]
) -> bool:

    """
    Sets up the websocket for Zerodha's Kite platform.
    :param auth_token: The auth-token model to use to se the feed up.
    :param total_instruments: The list of dicts that describe the instruments we need to subscribe to.
    :return: True if successful, else False.
    """

    # Declare the required global variables:
    global ZERODHA_INSTRUMENT_TOKENS
    global ZERODHA_INSTRUMENT_LOOKUP

    # Filter the input instruments. This has to be tuned to match what you get from the API.
    valid_broker_symbols = []
    valid_broker_tokens = []
    for i in total_instruments:
        symbol = i["symbol"]
        broker_token = i["broker_token"]
        if broker_token is not None and i["source"] == "zerodha":
            valid_broker_symbols.append(symbol)
            valid_broker_tokens.append(str(broker_token))

    # Debugging print:
    printer(len(total_instruments), len(valid_broker_symbols), len(valid_broker_tokens))

    # Create an instance of Zerodha's Kite connection:
    kite = KiteConnect(api_key = auth_token.auth["apiKey"])
    kite.set_access_token(auth_token.token["accessToken"])

    # Get the entire list of instruments:
    instruments = []
    instruments += kite.instruments(exchange = "NSE")
    instruments += kite.instruments(exchange = "NFO")
    instruments += kite.instruments(exchange = "BSE")
    instruments += kite.instruments(exchange = "BFO")
    instruments += kite.instruments(exchange = "MCX")
    instruments += kite.instruments(exchange = "CDS")
    instruments += kite.instruments(exchange = "BCD")

    # Convert the loaded instruments to their modelled form:
    instruments = [TradingSymbol.from_zerodha_kite(i) for i in instruments]

    # Create the lookup:
    for i in instruments:
        if str(i.brokerToken) in valid_broker_tokens:
            ZERODHA_INSTRUMENT_TOKENS.append(i.brokerToken)
            ZERODHA_INSTRUMENT_LOOKUP[i.brokerToken] = i.model_dump()
    printer("Zerodha instruments loaded.")

    # If there are no symbols or too many symbols, we return with failure:
    if not (1 <= len(ZERODHA_INSTRUMENT_TOKENS) <= 3_000):
        printer("Zerodha must have min. 1 and max. 3,000 symbols in one WS.", len(ZERODHA_INSTRUMENT_TOKENS))
        return False

    # Create the websocket to Zerodha:
    kite_ws = KiteTicker(
        api_key = auth_token.auth["apiKey"],
        access_token = auth_token.token["accessToken"]
    )

    # Assign the callbacks:
    kite_ws.on_connect = on_zerodha_kite_connect
    kite_ws.on_ticks = on_zerodha_kite_ticks

    # Run the websocket in a background thread and release the main thread:
    kite_ws.connect(threaded = True)

    # Done here:
    return True


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


async def get_auth_token(
        script_cred: dict,
        token_key: str,
        debug: bool
) -> CoreAuthTokenModel | None:

    """
    We need the full auth-token model to run the tick-ingestion script. This is a separate special coroutine that will
    use the existing async code to fetch the full document from the database.
    :param script_cred: The credentials of the script (needed to connect to the database).
    :param token_key: The token 'key' whose full auth-token model needs to be fetched.
    :param debug: Whether, or not, you would like to print debugging details.
    :return: The full auth-token of the account if found, or None if not found.
    """

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

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

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

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

    # ┏┓         ┓┓       ┓┏┏┓
    # ┃ ┏┓┏┓╋┏┓┏┓┃┃┏┓┏┓┏  ┃┃┏┛
    # ┗┛┗┛┛┗┗┛ ┗┛┗┗┗ ┛ ┛  ┗┛┗━

    controller = AllTradingController(
        cache = general_cache,
        http_client = httpx.AsyncClient(timeout = 10.0),
        alert_url = SCRIPT_DATA["alerts"]["url"],
        debug = debug
    )

    # Fetch and return the auth-token model for the given key:
    return await controller.get_token_from_key(
        mongo_data_conn = data_mongo,
        token_key = token_key
    )


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


def init(
        script_id: str,
        token_key: str,
        debug: bool
) -> 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 token_key: The key to identify the auth-token that must be used for connecting to the data feed.
    :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 TOKEN_KEY
    global AUTH_TOKEN
    global SCRIPT_DATA
    global kafka_producer
    global kafka_consumer

    # Basic stuff:
    TOKEN_KEY = token_key
    if not debug: printer.disable()

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

    # Get the script credentials:
    response = 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 = 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 producer that will produce tick-by-tick data that it receives from the stockbroker's server:
    producer_creds = script_cred["kafka"]["producer"]
    kafka_producer = ProducerKafka(
        config = ProducerKafka.create_config(
            bootstrap_servers = producer_creds["config"]["bootstrapServers"],
            security_protocol = producer_creds["config"].get("securityProtocol", "PLAINTEXT"),
            ca_file = producer_creds["config"].get("caFile"),
            cert_file = producer_creds["config"].get("certFile"),
            key_file = producer_creds["config"].get("keyFile"),
            client_id = f"{SERVER_HOSTNAME}_{TOKEN_KEY}",
            acks = producer_creds["config"].get("acks", 1),
            retries = producer_creds["config"].get("retries", 1),
            linger_ms = producer_creds["config"].get("lingerMs", 0),
            misc_json = {
                "queue.buffering.max.messages": 2_00_000
            }
        ),
        topic = producer_creds["topic"],
        serializer = JSONSerializer(),
        debug = debug
    )
    if not kafka_producer.connect():
        print("FATAL: KAFKA PRODUCER NOT CREATED!")
        return False
    printer("Kafka producer ready.")

    # Create the consumer that will listen to changes in watchlist:
    consumer_creds = script_cred["kafka"]["consumer"]
    kafka_consumer = ConsumerKafka(
        config = ConsumerKafka.create_config(
            bootstrap_servers = consumer_creds["config"]["bootstrapServers"],
            group_id = f"{SERVER_HOSTNAME}_{TOKEN_KEY}",
            auto_offset_reset = consumer_creds["config"].get("autoOffsetReset", "latest"),
            security_protocol = consumer_creds["config"].get("securityProtocol", "PLAINTEXT"),
            ca_file = consumer_creds["config"].get("caFile"),
            cert_file = consumer_creds["config"].get("certFile"),
            key_file = consumer_creds["config"].get("keyFile"),
            client_id = f"{SERVER_HOSTNAME}_{TOKEN_KEY}"
        ),
        topic = consumer_creds["topic"],
        serializer = JSONSerializer(),
        debug = debug
    )
    if not kafka_consumer.connect():
        print("FATAL: KAFKA CONSUMER NOT CREATED!")
        return False
    printer("Kafka consumer ready.")

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

    AUTH_TOKEN = asyncio.run(get_auth_token(
        script_cred = script_cred,
        token_key = token_key,
        debug = debug
    ))
    if not AUTH_TOKEN:
        print("FATAL: AUTH-TOKEN FETCHING FAILED!")
        return False
    printer("Auth-token fetched.")

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

    # First we load the tokens of interest that we need to monitor:
    response = http_client.post(
        url = r"https://api.thecaoffice.com/markets/watchlist/distincts",
        json = {"tokenKey": token_key}
    )
    instruments_of_interest = response.json()["data"]["rs0"]

    # Set the feed up by the stockbroker:
    live_feed = False
    if AUTH_TOKEN.client == "zerodhaKite":
        live_feed = setup_zerodha_kite_feed(
            auth_token = AUTH_TOKEN,
            total_instruments = instruments_of_interest
        )
    else: print(f"FATAL: INVALID TRADING CLIENT ({AUTH_TOKEN.client})!")

    # If our live feed setup failed:
    if not live_feed:
        print("FATAL: LIVE FEED SETUP FAILED!")
        return False

    printer("Live feed ready.")

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

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


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


def main():

    while True:
        pass


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


if __name__ == "__main__":

    # To get args. from the terminal:
    import argparse

    # Get the config. from the command-line:
    parser = argparse.ArgumentParser(
        description = (
            "Create one stateful process that gets tick-by-tick updates from stock brokers "
            "and produces them on the common Kafka broker. This process will be dedicated to one account."
        )
    )
    parser.add_argument(
        "-s", "--script-id",
        dest = "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",
        dest = "debug",
        action = "store_true",
        help = "Whether, or not, you want to see debugging messages in the terminal.",
        default = False
    )
    args = parser.parse_args()

    # Run the main script:
    if init(
        script_id = args.script_id,
        token_key = args.token_key,
        debug = args.debug
    ): main()
