"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Wednesday, 1st Jan., 2025.

    OBJECTIVE:

        To capture end-of-day market data and store it in the database.

    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_mysql_v2 import AsyncMySQL

# To make HTTP calls:
import httpx
import socket

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

# To work with NSE:
from utils_v2.nse.controllers.bhavcopy.equities import NSEEquitiesBhavCopy
from utils_v2.nse.controllers.bhavcopy.fno import NSEFNOBhavCopy

# To work with datatypes:
from typing import List

# For handling NaN values:
import pandas as pd

# For debugging:
from icecream import IceCreamDebugger

# For asynchronous operations:
import asyncio


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


# 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 = 120.0  # ..... Time to wait for receiving data.
    )
)

# For debugging:
printer = IceCreamDebugger(prefix = "EoD Data | ", includeContext = True)
no_context_printer = IceCreamDebugger(prefix = "EoD Data | ", includeContext = False)


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


# This script's config.:
SERVER_HOSTNAME = str(socket.gethostname())
SCRIPT_DATA = {}

# For the database:
sql_writer: AsyncMySQL | None = None

# To construct a report:
failed_to_parse_dates = []
failed_eq_dates = []
failed_fo_dates = []


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


async def init(
        script_id: 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 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 sql_writer

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

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

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

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

    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 CONNECTION FAILED!")
        return False

    printer("MariaDB connected.")

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

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


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


async def capture_eq_eod_data(
        nse_conn: NSEEquitiesBhavCopy,
        target_date: datetime.datetime
) -> bool:

    # Start by assuming failure:
    success = False

    # Get the data:
    no_context_printer("Fetching EQ data.")
    nse_response = await nse_conn.get_data(
        target_date = target_date,
        return_raw = False,
        retry_count = 3,
        backoff_seconds = 0.5,
        backoff_multiplier = 1.1
    )
    if not nse_response.success:
        no_context_printer(nse_response.method, nse_response.url, nse_response.httpCode)
        return success

    # Build the query and data content:
    query_str = (
        "INSERT INTO `eod_market_data_temp` (exchange, symbol, segment, type, underlying, is_index, expiry, strike, "
        "prev_close, open, high, low, close, ltp, vwap, tot_vol, tot_cash, delivery_vol, delivery_pct, oi, oi_chg, "
        "date, ts, tz, scrape_ts) "
        "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);"
    )
    query_data = []
    for data in nse_response.data:
        one_query_data = (
            "nse",  # ..................................... exchange
            data["symbol"],  # ............................ symbol
            data["segment"],  # ........................... segment
            None,  # ...................................... type
            None,  # ...................................... underlying
            False,  # ..................................... is_index
            None,  # ...................................... expiry
            None,  # ...................................... strike
            data["prevClose"],  # ......................... prev_close
            data["open"],  # .............................. open
            data["high"],  # .............................. high
            data["low"],  # ............................... low
            data["close"],  # ............................. close
            data["ltp"],  # ............................... ltp
            data["vwap"],  # .............................. vwap
            data["totVol"],  # ............................ tot_vol
            data["totCash"],  # ........................... tot_cash
            data["deliveryVol"],  # ....................... delivery_vol
            data["deliveryPct"],  # ....................... delivery_pct
            None,  # ...................................... oi
            None,  # ...................................... oi_chg
            target_date,  # ............................... date
            data["ts"].replace(tzinfo = None),  # ......... ts
            data["tz"],  # ................................ tz
            data["scrapeTs"].replace(tzinfo = None),  # ... scrape_ts
        )
        one_query_data = [None if pd.isna(d) else d for d in one_query_data]
        query_data.append(one_query_data)

    # Run the commands:
    no_context_printer("Saving EQ data to DB.")
    rows_affected, db_response = await sql_writer.execute_many(
        query = query_str,
        data = query_data
    )

    # Done here:
    success = True if rows_affected else False
    no_context_printer(success)
    return success


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


async def capture_fo_eod_data(
        nse_conn: NSEFNOBhavCopy,
        target_date: datetime.datetime
) -> bool:

    # Start by assuming failure:
    success = False

    # Get the data:
    no_context_printer("Fetching FO data.")
    nse_response = await nse_conn.get_data(
        target_date = target_date,
        return_raw = False,
        retry_count = 3,
        backoff_seconds = 0.5,
        backoff_multiplier = 1.1
    )
    if not nse_response.success:
        no_context_printer(nse_response.method, nse_response.url, nse_response.httpCode)
        return success

    # Build the query and data content:
    query_str = (
        "INSERT INTO `eod_market_data_temp` (exchange, symbol, segment, type, underlying, is_index, expiry, strike, "
        "prev_close, open, high, low, close, ltp, vwap, tot_vol, tot_cash, delivery_vol, delivery_pct, oi, oi_chg, "
        "date, ts, tz, scrape_ts) "
        "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);"
    )
    query_data = []
    for data in nse_response.data:
        expiry = data["expiryTs"]
        if isinstance(expiry, pd._libs.tslibs.timestamps.Timestamp): expiry = expiry.to_pydatetime()
        if isinstance(expiry, datetime.datetime):
            expiry = date_time.to_timezone(
                datetime_object = expiry,
                timezone = date_time.TIMEZONE_UTC
            ).strftime("%Y-%m-%d")
        one_query_data = (
            "nse",  # ..................................... exchange
            data["symbol"],  # ............................ symbol
            data["segment"],  # ........................... segment
            data["type"],  # .............................. type
            data["underlying"],  # ........................ underlying
            data["isIndex"],  # ........................... is_index
            expiry,  # .................................... expiry
            data["strike"],  # ............................ strike
            data["prevClose"],  # ......................... prev_close
            data["open"],  # .............................. open
            data["high"],  # .............................. high
            data["low"],  # ............................... low
            data["close"],  # ............................. close
            None,  # ...................................... ltp
            None,  # ...................................... vwap
            data["totVol"],  # ............................ tot_vol
            data["totCash"],  # ........................... tot_cash
            None,  # ...................................... delivery_vol
            None,  # ...................................... delivery_pct
            data["oi"],  # ................................ oi
            data["oiChg"],  # ............................. oi_chg
            target_date,  # ............................... date
            data["ts"].replace(tzinfo = None),  # ......... ts
            data["tz"],  # ................................ tz
            data["scrapeTs"].replace(tzinfo = None),  # ... scrape_ts
        )
        one_query_data = [None if pd.isna(d) else d for d in one_query_data]
        query_data.append(one_query_data)

    # Run the commands:
    no_context_printer("Saving FO data to DB.")
    rows_affected, db_response = await sql_writer.execute_many(
        query = query_str,
        data = query_data
    )

    # Done here:
    success = True if rows_affected else False
    no_context_printer(success)
    return success


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


async def main(
        target_dates: List[str],
        interval: float = 5.0,
        capture_eq: bool = True,
        capture_fo: bool = True
):

    # Declare the needed global variables:
    global failed_to_parse_dates
    global failed_eq_dates
    global failed_fo_dates

    # Create the NSE connection instances:
    nse_eq_bhavcopy = NSEEquitiesBhavCopy(http_client = http_client, debug = False)
    nse_fo_bhavcopy = NSEFNOBhavCopy(http_client = http_client, debug = False)

    # Load the data for each date:
    printer(len(target_dates))
    for dt_str in target_dates:

        # Try parsing the date:
        dt = date_time.parse_date_time(
            input_value = dt_str,
            date_formats = ["%Y-%m-%d", "%Y%m%d"]
        )

        # If we failed to parse the date:
        if dt is None:
            no_context_printer("FAILED to parse date.", dt_str)
            failed_to_parse_dates.append(dt_str)
            continue

        no_context_printer("Picked date.", dt_str)

        # Fetch the equity data:
        if capture_eq:
            success = await capture_eq_eod_data(
                nse_conn = nse_eq_bhavcopy,
                target_date = dt
            )
            if not success:
                no_context_printer("FAILED equity for date.", dt_str)
                failed_eq_dates.append(dt_str)

        # Fetch the derivative data:
        if capture_fo:
            success = await capture_fo_eod_data(
                nse_conn = nse_fo_bhavcopy,
                target_date = dt
            )
            if not success:
                no_context_printer("FAILED derivatives for date.", dt_str)
                failed_fo_dates.append(dt_str)

        # Pause for a while to not get rate-limited/blocked:
        await asyncio.sleep(interval)

    # Show the report:
    printer("Process complete.", failed_to_parse_dates, capture_eq, failed_eq_dates, capture_fo, failed_fo_dates)


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


if __name__ == "__main__":

    # To get args. from the terminal:
    import argparse

    # Get the config. from the command-line:
    parser = argparse.ArgumentParser(
        description = (
            "To fetch and store EoD data for one or more dates from NSE's BhavCopy section. "
            "The data will be fetched for both, equities and derivatives."
        )
    )
    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", "--dates",
        dest = "dates",
        type = str,
        help = "One or more comma-separated dates in 'YYYY-MM-DD' format for which you would like to get the EoD data.",
        default = None
    )
    parser.add_argument(
        "-i", "--interval",
        dest = "interval",
        type = float,
        help = "The delay, in seconds, between processing two dates.",
        default = 2.5
    )
    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
    )
    parser.add_argument(
        "-e", "--eq",
        dest = "eq",
        action = "store_true",
        help = "To enable equity data capture.",
        default = False
    )
    parser.add_argument(
        "-f", "--fo",
        dest = "fo",
        action = "store_true",
        help = "To enable derivatives data capture.",
        default = False
    )
    args = parser.parse_args()


    async def runner():

        # Prepare the list of target dates:
        if args.dates is None: target_dates = [datetime.datetime.now().strftime("%Y-%m-%d")]
        else: target_dates = [d.strip() for d in args.dates.split(",")]

        # Initialize and run the main code:
        if await init(
            script_id = args.script_id,
            debug = args.debug
        ): await main(
            target_dates = target_dates,
            interval = args.interval,
            capture_eq = args.eq,
            capture_fo = args.fo
        )

        # Disconnect from the database:
        disconnected = await sql_writer.disconnect()


    asyncio.run(runner())
