"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 17th Jan., 2025.

    OBJECTIVE:

        To automatically synchronize the mails by fetching data from the user's third-party mail client's server to your
        database. This is done in batches of auth-tokens.

    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
import socket

# 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.database.async_mysql_v2 import AsyncMySQL
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

# Controllers:
from controllers_v2.message.mail.all_mail import AllMailController
from controllers_v2.message.mail.gmail import GmailController
from controllers_v2.core.ai.llm import CoreLLMController

# Mail clients:
from utils_v2.goog.controllers.gmail.gmail_client import AsyncGmailClient

# To make HTTP calls:
import httpx

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

# Models:
from models.core.auth_token import CoreAuthTokenModel
from models.core.user import CoreUserInfoModel
from models.message.mail.sync import MailSyncOneResult, MailSyncManyResults
from models.core.ai.llm import LLMInput, LLMInputMessage

# For asynchronous activities:
import asyncio

# To work with various datatypes:
from typing import List

# For random values:
import random

# Debugging:
from icecream import IceCreamDebugger


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


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


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

# For databases:
data_mongo: AsyncMongo | None = None
sql_writer: AsyncMySQL | None = None

# For caching:
redis_cache: AsyncRedisCache | None = None

# Controllers:
mail_controller: AllMailController | None = None
gmail_controller: GmailController | None = None
llm_controller: CoreLLMController | None = None

# Mail clients:
gmail_client: AsyncGmailClient | None = None


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


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 sql_writer
    global redis_cache
    global mail_controller
    global gmail_controller
    global llm_controller
    global gmail_client

    # 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 NOT CONNECTED!")
        return False
    no_context_printer("MariaDB ready.")

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

    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
    )
    if not await data_mongo.connect():
        print("FATAL: MONGO-DB NOT CONNECTED!")
        return False
    no_context_printer("MongoDB ready.")

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

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

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

    # Messages / Mail Controllers:
    mail_controller = AllMailController(
        cache = redis_cache,
        http_client = http_client,
        alert_url = SCRIPT_DATA["alerts"]["url"],
        debug = False
    )
    gmail_controller = GmailController(
        cache = redis_cache,
        http_client = http_client,
        alert_url = SCRIPT_DATA["alerts"]["url"],
        debug = False
    )

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

    # Create an instance to handle GMail-related activities:
    gmail_client = AsyncGmailClient(
        service_name = "gmail",
        oauth_json = script_cred["google"]["oauth"]["tcaoff"],
        http_client = http_client,
        redirect_url = r"https://api.thecaoffice.com/converse/mail/callback/gmail",
        debug = False,
        debug_prefix = "Gmail (M) | ",
        debug_only_errors = True
    )

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

    # For LLMs:
    llm_controller = CoreLLMController(
        llm_creds = {
            "model": script_cred["openAi"]["model"],
            "openai_api_key": script_cred["openAi"]["openai_api_key"]
        },
        cache = redis_cache,
        alert_url = SCRIPT_DATA["alerts"]["url"],
        http_client = http_client,
        debug = debug,
        debug_prefix = "AI (LLM) | ",
        debug_only_errors = True
    )

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

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


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


async def sync_one_account(auth_token: CoreAuthTokenModel):

    """
    To sync one third-party mail account.
    :param auth_token: The auth-token that gives access to the mail account.
    :return: ??
    """

    # Note down the time at which the attempt to sync the account is being made:
    now = date_time.get_current_utc_date_time(as_string = False)

    # Start by assuming failure:
    client_controller = None
    client_connector = None
    sync_results = MailSyncManyResults()

    # Figure out which mail client has to be used:
    match auth_token.client:
        case "gmail": client_controller, client_connector = gmail_controller, gmail_client
        case _: client_controller, client_connector = None, None

    # Try to sync mails:
    if client_controller is not None and client_connector is not None:
        try: sync_results = await client_controller.sync_mails(
            sql_conn = sql_writer,
            mongo_data_conn = data_mongo,
            mail_client = client_connector,
            auth_token = auth_token,
            user_info = None,
            llm = llm_controller,
            force_sync = False,
            start_date = now - datetime.timedelta(days = 1),
            end_date = now,
            max_count = 100,
            session_token = None
        )
        except Exception as exception:
            pass
            # printer(exception)
    # else: printer("Invalid/unimplemented client.", auth_token.client)

    # Release the auth-token from the batch:
    await mail_controller.release_token_from_batch_by_id(
        mongo_data_conn = data_mongo,
        token_id = auth_token.authTokenId,
        sync_after_ts = now + datetime.timedelta(seconds = auth_token.syncFreq or 300),
        last_sync_ts = now
    )

    # Done here:
    if sync_results.totalCount:
        printer(
            auth_token.clientUserId["email"],
            sync_results.totalCount,
            sync_results.newCount,
            sync_results.attemptedCount,
            sync_results.successCount,
            sync_results.failureCount
        )


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


async def sync_accounts(batch_size: int) -> None:

    """
    To go into an indefinite loop and keep sync'ing mails for many accounts.
    :param batch_size: The no. of accounts to pick in every batch.
    :return: None.
    """

    while True:

        # Get batches of auth-tokens to work with:
        auth_tokens_batch = await mail_controller.get_batches_to_sync(
            mongo_data_conn = data_mongo,
            limit = batch_size,
            batch_timeout_seconds = 10
        )
        no_context_printer(len(auth_tokens_batch))

        # Process each batch:
        now = date_time.get_current_utc_date_time(as_string = False)
        tasks = [sync_one_account(auth_token = auth_token) for auth_token in auth_tokens_batch]
        results = await asyncio.gather(*tasks)

        # Small delay to not overload the database:
        await asyncio.sleep(100.0 + (random.random() * 20.0))
        # if len(auth_tokens_batch) == 0: await asyncio.sleep(4.0 + (random.random() * 2.0))
        # else: await asyncio.sleep(1.0 + (random.random() * 2.0))


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


if __name__ == "__main__":

    # To get args from the terminal:
    import argparse

    # Get the config from the command-line:
    parser = argparse.ArgumentParser(description = f"To automatically sync. the mails for all users.")
    parser.add_argument(
        "-b", "--batch-size",
        type = int,
        help = "How many auth-tokens to load at once to sync.",
        default = 25
    )
    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()

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

    async def main():
        if await init(
            script_id = args.script_id,
            debug = args.debug
        ): await sync_accounts(
            batch_size = args.batch_size
        )

    asyncio.run(main())
