"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 24th Jan., 2025.

    OBJECTIVE:

        Live market ticks are broadcasted over Kafka. Here, we try to capture them and use them to run strategies.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


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


# To make sibling directories accessible for imports:
import sys

import pandas as pd

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_mysql_v2 import AsyncMySQL
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 scheduling and cron:
from scheduler.asyncio import Scheduler

# For MongoDB:
from bson.objectid import ObjectId

# Debugging:
from icecream import IceCreamDebugger


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


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

# For cache:
redis_cache: AsyncRedisCache | None = None

# For database(s):
sql_writer: AsyncMySQL | None = None
sql_reader: AsyncMySQL | None = None
data_mongo: AsyncMongo | None = None

# For kafka:
kafka_consumer: ConsumerKafka | None = None

# Session-awareness and maintenance of this script's state:
SCRIPT_DATA = {}
STRATEGY_REFERENCE = {}
ACTIVE_CALLS = {}
strategy_reference_lock = asyncio.Semaphore(1)
ticks_since_telegram = 0


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


async def mark_call_as_active(
        tick_ref: dict,
        tick: dict
) -> bool:

    global ACTIVE_CALLS

    # get the keys:
    tick_key = tick_ref["tickKey"]
    redis_key = tick_ref["redisKey"]

    # Mark locally:
    # print(f"CALL ({redis_key}):", json.to_string(tick_ref))
    ACTIVE_CALLS[redis_key] = tick_ref

    # Mark in cache:
    await redis_cache.set(
        key = redis_key,
        value = tick_ref
    )

    # Mark in SQL:
    proc_name = "strategy_trade_add"
    proc_args = (
        tick_ref["userId"],  # ........................................................ p_user_id
        tick_ref["billingAccountId"],  # .............................................. p_billing_account_id
        0,  # ......................................................................... p_entity_integration_id
        date_time.get_current_ist_date_time().strftime("%Y-%m-%d"),  # ................ p_date
        tick_ref["exchange"],  # ...................................................... p_exchange
        tick_ref["segment"],  # ....................................................... p_segment
        tick_ref["symbol"],  # ........................................................ p_symbol
        tick_ref["expiry"],  # ........................................................ p_expiry
        tick_ref["right"],  # ......................................................... p_right
        tick_ref["strike"],  # ........................................................ p_strike
        0,  # ......................................................................... p_qty
        0,  # ......................................................................... p_price
        tick["ltp"],  # ............................................................... p_refPrice
        tick_ref["eqQty"] if tick_ref["segment"] == "EQ" else tick_ref["foQty"],  # ... p_refQty
        tick_ref["exchangeCode"]  # .................................................. p_exchange_code
    )
    db_json, db_exception = await sql_reader.call_procedure_and_get_json(
        procedure_name = proc_name,
        procedure_args = proc_args,
        return_exception = True,
        retry_count = 3,
        backoff_seconds = 0.5,
        backoff_multiplier = 1.1
    )

    # Alert if the portfolio adding fails:
    if db_json["status"] != 1 or db_exception:
        await send_telegram(
            message = (
                "*Strategies (0) SQL Procedure Failure*\n\n"
                f"Procedure Name: `{proc_name}`\n\n"
                f"Procedure Args: `{proc_args}`\n\n"
                f"Message: `{db_json['message']}`\n\n"
                f"Exception: `{db_exception}`\n\n"
            ),
            message_type = "error"
        )

    # Done here:
    return True


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


async def call_is_active(
        tick_ref: dict,
        tick: dict
) -> bool:

    global ACTIVE_CALLS

    # get the keys:
    tick_key = tick_ref["tickKey"]
    redis_key = tick_ref["redisKey"]

    # Check if the call is active:
    if ACTIVE_CALLS.get(redis_key): return True
    if await redis_cache.get(redis_key):
        ACTIVE_CALLS[redis_key] = tick_ref
        return True
    return False


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


def key_from_tick(tick: dict) -> str:

    if tick["broker"] == "zerodhaKite":

        is_eq = True if tick["type"] == "EQ" else False

        exchange = {
            "NFO": "NSE",
        }.get(tick["exchange"], tick["exchange"])

        segment = {
            "NFO-FUT": "FUT",
            "BFO-FUT": "FUT",
            "MCX-FUT": "FUT",
        }.get(tick["segment"], tick["segment"])

        symbol = tick["symbol"]
        name = tick["name"]

        expiry = tick["expiry"]

        if is_eq: key_items = [exchange, tick["type"], symbol]
        else: key_items = [exchange, segment, name, expiry]
        return "_".join(key_items)


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


async def get_reference_for_tick(tick_key: str, tick: dict) -> dict | None:

    tick_reference = None

    async with strategy_reference_lock:
        for ref in STRATEGY_REFERENCE:
            if ref["tickKey"] == tick_key:
                tick_reference = ref
                break

    return tick_reference


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


async def test_one_tick(tick: dict) -> None:

    # Get the reference for this ticker:
    tick_key = key_from_tick(tick)
    tick_reference = await get_reference_for_tick(tick_key, tick)
    if not tick_reference: return
    redis_key = tick_reference["redisKey"]

    # The call must NOT be active already:
    # print(f"FOUND REF ({tick_key} | {tick['ltp']})", json.to_string(tick_reference, default = str))
    if await call_is_active(
        tick_ref = tick_reference,
        tick = tick
    ): return

    # Check if the price is favourable:
    if (
            (
                    tick_reference["entry"] == "BUY" and
                    tick["ltp"] <= tick_reference["rate"]
            ) or
            (
                    tick_reference["entry"] == "SELL" and
                    tick["ltp"] >= tick_reference["rate"]
            )
    ):
        mark_success = await mark_call_as_active(
            tick_ref = tick_reference,
            tick = tick
        )
        if not mark_success: await send_telegram(
            message = (
                "*Strategies (0) Alert*\n\n"
                "Message: `Failed to mark active call.`\n\n"
                f"Tick Key; `{tick_key}`\n"
                f"Redis Key; `{redis_key}`\n"
                f"Entry: `{tick_reference['entry']}`\n"
                f"Rate: `{tick_reference['rate']}`\n"
                f"LTP: `{tick['ltp']}`"
            ),
            message_type = "error"
        )
        else: await send_telegram(
            message = (
                "*Strategies (0) Alert*\n\n"
                "Message: `Marked active call.`\n\n"
                f"Tick Key; `{tick_key}`\n"
                f"Redis Key; `{redis_key}`\n"
                f"Entry: `{tick_reference['entry']}`\n"
                f"Rate: `{tick_reference['rate']}`\n"
                f"LTP: `{tick['ltp']}`"
            ),
            message_type = "info"
        )


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


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

    # declare the needed global variables:
    global ticks_since_telegram

    # Test needed strategies here:
    # print("TICK 0:", json.to_string(ticks[0]))
    tasks = [test_one_tick(t) for t in ticks]
    results = await asyncio.gather(*tasks)

    # 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"*Strategies (0) Alert*\nProcessed {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
    """

    no_context_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: 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 test_strategies(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 redis_cache
    global sql_writer
    global sql_reader
    global kafka_consumer

    # Basic stuff:
    if debug: printer.enable()
    no_context_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:
    no_context_printer("Cred and Data loaded.")

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

    sql_writer = AsyncMySQL(
        pool_size = script_cred["mariaDb"]["write"]["poolSize"],
        host = script_cred["mariaDb"]["write"]["host"],
        user = script_cred["mariaDb"]["write"]["user"],
        password = script_cred["mariaDb"]["write"]["password"],
        database = script_cred["mariaDb"]["write"]["database"]
    )
    if not await sql_writer.connect():
        print("FATAL: MARIA-DB WRITER CONNECTION FAILED!")
        return False

    sql_reader = AsyncMySQL(
        pool_size = script_cred["mariaDb"]["read"]["poolSize"],
        host = script_cred["mariaDb"]["read"]["host"],
        user = script_cred["mariaDb"]["read"]["user"],
        password = script_cred["mariaDb"]["read"]["password"],
        # database = script_cred["mariaDb"]["read"]["database"]
        database = "caOffice"
    )
    if not await sql_reader.connect():
        print("FATAL: MARIA-DB READER CONNECTION FAILED!")
        return False

    no_context_printer("MariaDB connected.")

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

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

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

    # Caching connections:
    redis_cache = AsyncRedisCache(
        connection_string = script_cred["redisCache"]["strategies"]["connectionString"],
        debug = False,
        debug_prefix = "Cache | "
    )
    if not await redis_cache.connect():
        print("FATAL: REDIS CACHE NOT CREATED!")
        return False
    no_context_printer("Redis cache ready.")

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

    # If everything went well, we return with success:
    no_context_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 refresh_strategy_reference() -> bool:

    # Declare the required global variables:
    global STRATEGY_REFERENCE

    no_context_printer("Refreshing strategy ref.")

    # Query the database:
    proc_name = "strategy_results_all"
    # proc_args = (0, 184, 0)
    proc_args = ()
    db_json, db_exception = await sql_reader.call_procedure_and_get_json(
        procedure_name = proc_name,
        procedure_args = proc_args,
        return_exception = True,
        retry_count = 3,
        backoff_seconds = 0.5,
        backoff_multiplier = 1.1
    )

    # If the database call fails:
    if db_json["status"] != 1 or db_exception:
        await send_telegram(
            message = (
                "*Strategies (0) SQL Procedure Failure*\n\n"
                f"Procedure Name: `{proc_name}`\n\n"
                f"Procedure Args: `{proc_args}`\n\n"
                f"Message: `{db_json['message']}`\n\n"
                f"Exception: `{db_exception}`\n\n"
            ),
            message_type = "error"
        )
        return False

    # Construct the reference structure:
    formatted_reference = []
    for i in db_json["data"]["rs0"]:
        key_items = [i["exchange"], i["segment"], i["exchange_code"]]
        if i["expiry_date"]: key_items.append(str(i["expiry_date"]))
        if i["right"]:
            key_items.append(str(i["right"]))
            key_items.append(str(i["strike"]))
        tick_key = "_".join(key_items)
        strategy_id = i["strategy_id"]
        strategy_name = i["name"]
        strategy_entry = i["entry"].upper().strip()
        formatted_reference.append({
            "userId": i["user_id"],
            "billingAccountId": i["billing_account_id"],
            "eqQty": i["eq_qty"],
            "foQty": i["fno_qty"],
            "tickKey": tick_key,
            "redisKey": f"{tick_key}_{strategy_name}_{strategy_id}_{strategy_entry}",
            "stratId": i["strategy_id"],  # .......... As inserted by the user.
            "stratName": strategy_name,  # .......... As inserted by the user.
            "watchlistId": i["watch_list_id"],  # ... A unique id in case 2 users come up with the same name.
            "entry": strategy_entry,  # ............. BUY/SELL.
            "rate": i["entry_rate"],  # ............. The price at which to enter the trade.
            "t": i["target1"],  # ................... The 1st target.
            "sl": i["stop_loss"],  # ................. The stop loss value.
            "exchange": i["exchange"],
            "segment": i["segment"],
            "expiry": i["expiry_date"],
            "strike": i["strike"],
            "right": i["right"],
            "symbol": i["symbol"],
            "exchangeCode": i["exchange_code"],
        })

    # Save the new reference in the global variable:
    async with strategy_reference_lock:
        STRATEGY_REFERENCE = formatted_reference
        # print(json.to_string(list(STRATEGY_REFERENCE)))

    # Done here:
    return True


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


async def heartbeat(interval_seconds: float = 300) -> None:

    """
    Just to let the admins know that the system is up and running.
    :param interval_seconds: The amount of time after which the heartbeat must be sent.
    :return: None.
    """

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


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


async def main(
        start_time: datetime.datetime,
        end_time: datetime.datetime,
        proc_interval: int,
        heartbeat_interval: int = 1_800,
        debug: bool = False
):

    # Get the reference before the scheduler kicks in:
    await refresh_strategy_reference()

    # Start configuring the scheduler:
    no_context_printer("Configuring the schedule-manager.")
    schedule_manager = Scheduler()

    # Create all the timestamps at which the job must be done:
    all_job_ts = []
    offset_seconds = 0
    while True:
        ts = start_time + datetime.timedelta(seconds = offset_seconds)
        if ts > end_time: break
        all_job_ts.append(ts.time())
        offset_seconds += proc_interval

    # Add the jobs:
    for ts in all_job_ts: schedule_manager.daily(ts, refresh_strategy_reference)
    no_context_printer(len(all_job_ts))

    # Run the heartbeat task and the infinite tick-reading loop:
    tasks = [
        ticks_from_kafka(
            consumer = kafka_consumer,
            fetch_count = 500,
            fetch_timeout = 2.5
        ),
        # heartbeat(interval_seconds = heartbeat_interval)
    ]
    await asyncio.gather(*tasks)


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


if __name__ == "__main__":

    no_context_printer("Main.")

    # To get args from the terminal:
    import argparse

    # Get the config from the command-line:
    parser = argparse.ArgumentParser(description = f"To implement trading strategies (0) for Mr. M. Bhandari.")
    parser.add_argument(
        "-s", "--script-id",
        type = str,
        help = "The id of this script (will affect the loaded config)."
    )
    parser.add_argument(
        "--start-time",
        type = str,
        help = "The 24-hr time of the day (in 'HH:MM:SS' format) from which the data can be refreshed."
    )
    parser.add_argument(
        "--end-time",
        type = str,
        help = "The 24-hr time of the day (in 'HH:MM:SS' format) till which the data must be refreshed."
    )
    parser.add_argument(
        "--proc-interval",
        type = int,
        help = "The no. of seconds after which you would like to refresh the data available from the stored procedure.",
        default = 300
    )
    parser.add_argument(
        "--heartbeat-interval",
        type = int,
        help = "The no. of seconds after which you would like to send out a heartbeat to the admins.",
        default = 1_800
    )
    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()

        # Parse the inputs:
        start_time = datetime.datetime.strptime(args.start_time, "%H:%M:%S")
        end_time = datetime.datetime.strptime(args.end_time, "%H:%M:%S")

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

    asyncio.run(runner())
