"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 12th Dec., 2024

    OBJECTIVE:

        To handle all auth-tokens from one place.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""

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


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# For Quart:
from quart import current_app

# My async utils:
from utils_v2.string import json
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, AsyncMongoStorage

# Base model:
from controllers.base import BaseModel

# Data models:
from models.core.user import CoreUserInfoModel
from models.core.auth_token import CoreAuthTokenModel
from models.core.message import CoreMessageModel
from models.api.mail.sync import MailSyncOneResult, MailSyncManyResults

# Mail Clients:
from utils_v2.goog.gmail.gmail_client import AsyncGMailClient
from utils_v2.goog.models.auth_tokens import GoogleAuthTokens

# To work with MongoDB:
from bson import ObjectId
from pymongo import InsertOne, UpdateOne, ReplaceOne

# To work with LLMs:
from controllers.core.ai.llm import CoreLLMController
from models.core.ai.llm import LLMInput, LLMOutput

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

# To parse the HTML content in the mail:
from bs4 import BeautifulSoup

# To work with date and time:
import datetime

# For asynchronous activities:
import asyncio


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


# --- Nothing Yet


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


# --- Nothing Yet


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


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class MailController:

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

    # For AI Magic through LLMs:
    PROMPT_TEMPLATE = [
        {
            "role": "system",
            "content": (
                "You're a mail summary expert that summarizes mails in 150 chars or less. "
                "If available, show login info like username and OTPs in your summary."
                "If no login info is provided, please don't worry; just summarize what you see."
            )
        }
    ]

    # ┓┏  ┓
    # ┣┫┏┓┃┏┓┏┓┏┓┏
    # ┛┗┗ ┗┣┛┗ ┛ ┛
    #      ┛

    def extract_plaintext_parts(
            self,
            payload: dict
    ) -> List[str]:

        # Start with just a holder:
        text_parts = []

        # If a direct text/plain part occurs,
        # we just add it to the list:
        if (
                payload["contentMainType"] == "text" and
                payload["contentSubType"] == "plain"
        ):
            text_parts.append(payload["payload"])

        # If a direct text/html part occurs,
        # we just add it to the list:
        if (
                payload["contentMainType"] == "text" and
                payload["contentSubType"] == "html"
        ):
            html_parser = BeautifulSoup(payload["payload"], "html.parser")
            text_parts.append(html_parser.get_text())

        # If a multipart/alternative part occurs,
        # we pick just the ready plaintext part:
        if (
                payload["contentMainType"] == "multipart" and
                payload["contentSubType"] == "alternative"
        ):
            for part in payload["payload"]:
                if part["contentSubType"] == "plain":
                    text_parts.append(part["payload"])

        # If a multipart/mixed or multipart/related part occurs,
        # we use recursion to look for plaintext parts nested inside:
        if (
                payload["contentMainType"] == "multipart" and
                (
                        payload["contentSubType"] == "mixed" or
                        payload["contentSubType"] == "related"
                )
        ):
            for part in payload["payload"]:
                text_parts += self.extract_plaintext_parts(payload = part)

        # Done here:
        return text_parts

    async def summarize_mail_with_ai(
            self,
            mongo_conn: AsyncMongo,
            user_info: CoreUserInfoModel,
            llm: CoreLLMController,
            message: CoreMessageModel
    ) -> LLMOutput:

        # Extract the text from the message here:
        text_parts = self.extract_plaintext_parts(payload = message.message["payload"])
        text = "\n".join(text_parts)

        # Invoke the LLM and return the response:
        return await llm.invoke(
            mongo_conn = mongo_conn,
            user_info = user_info,
            llm_input = LLMInput(
                messages = self.PROMPT_TEMPLATE + [
                    {
                        "role": "human",
                        "content": f"Please summarize this mail: \"\"\"{text}\"\"\""
                    }
                ]
            )
        )

    def drop_attachments(
            self,
            payload: dict
    ):

        # If the part is some sort of file:
        if payload["contentMainType"] not in ["multipart", "text"]:
            payload["payload"] = None
            payload["payloadId"] = None
            payload["payloadUrl"] = None

        # If the payload is of multipart type,
        # we use recursion to look inside it:
        elif payload["contentMainType"] == "multipart":
            for part in payload["payload"]:
                self.drop_attachments(part)

        # Done here:
        return payload

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

    @staticmethod
    async def get_token_key(
            db_conn: AsyncMySQL,
            mongo_conn: AsyncMongo,
            auth_token: CoreAuthTokenModel,
            session_token: str = None
    ) -> ObjectId:

        # Simply call the core model:
        return await current_app.core_auth_token_controller.get_token_key(
            db_conn = db_conn,
            mongo_conn = mongo_conn,
            auth_token = auth_token,
            token_notes = {
                "email": None
            },
            session_token = session_token,
        )

    @staticmethod
    async def set_token(
            db_conn: AsyncMySQL,
            mongo_conn: AsyncMongo,
            token_key: ObjectId | str,
            auth_token: CoreAuthTokenModel,
            session_token: str = None
    ) -> bool:

        # Simply call the core model:
        return await current_app.core_auth_token_controller.set_token(
            db_conn = db_conn,
            mongo_conn = mongo_conn,
            token_key = token_key,
            auth_token = auth_token,
            token_notes = {
                "email": auth_token.token["email"],
                "displayName": auth_token.token.get("displayName"),
                "displayPictureUrl": auth_token.token.get("displayPictureUrl"),
            },
            session_token = session_token,
        )

    @staticmethod
    async def get_token(
            mongo_conn: AsyncMongo,
            token_key: ObjectId | str = None,
    ) -> CoreAuthTokenModel | None:

        # Simply call the core model:
        return await current_app.core_auth_token_controller.get_token(
            mongo_conn = mongo_conn,
            token_key = token_key
        )

    @staticmethod
    async def get_tokens(
            mongo_conn: AsyncMongo,
            token_keys: ObjectId | str = None,
    ) -> List[CoreAuthTokenModel] | None:

        # Simply call the core model:
        return await current_app.core_auth_token_controller.get_tokens(
            mongo_conn = mongo_conn,
            token_keys = token_keys
        )

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

    # In this section, we pull mails from the third-party clients (like GMail), and store them on our server. This makes
    # those mails available on the platform.

    async def __sync_one_gmail(
            self,
            mongo_conn: AsyncMongo,
            user_info: CoreUserInfoModel,
            auth_token: CoreAuthTokenModel,
            mail_client: AsyncGMailClient,
            google_tokens: GoogleAuthTokens,
            message_id: str,
            llm: CoreLLMController = None,
            force_sync: bool = False
    ) -> MailSyncOneResult:

        # Start by assuming failure:
        sync_result = MailSyncOneResult()

        # If we've not been forced to re-sync the mail message,
        # we first check if the mail already exists in our database:
        if not force_sync:
            mail_records = await current_app.core_message_controller.get_previews(
                mongo_conn = mongo_conn,
                token_ids = [ObjectId(auth_token.authTokenId)],
                limit = 1,
                skip = 0,
                additional_filter = {
                    "tokenId": ObjectId(auth_token.authTokenId),
                    "serviceType": auth_token.serviceType,
                    "client": auth_token.client,
                    "clientMessageId": message_id
                }
            )
            if mail_records:
                sync_result.success = True
                sync_result.message = (
                    f"gmail message '{message_id}' already "
                    f"sync'd on '{mail_records[0].syncTs} (UTC)'"
                )
                return sync_result

        # Now that we know that we have to fetch the mail from GMail:
        client_response = await mail_client.get_message(
            tokens = google_tokens,
            message_id = message_id,
            return_raw = False
        )

        # If we didn't get the mail from GMail;
        if not client_response.success:
            sync_result.message = f"gmail (messageId: '{message_id}'): {client_response.message}"
            return sync_result

        # HANDLE ATTACHMENTS HERE:
        client_response.data["payload"] = self.drop_attachments(client_response.data["payload"])

        # Now we structure the message into the model:
        mail_message = CoreMessageModel(
            ts = client_response.data["ts"],
            syncTs = date_time.get_current_utc_date_time(as_string = False),
            tokenId = auth_token.authTokenId,
            serviceType = auth_token.serviceType,
            client = auth_token.client,
            clientMessageId = message_id,
            clientThreadId = client_response.data["threadId"],
            isSent = False,
            isBroadcast = False,
            sentSuccessfully = False,
            sender = client_response.data["from"][0]["name"],
            chat = None,
            message = client_response.data,
            snippet = client_response.data["subject"],
            aiSnippet = None,
            tags = ["email", "gmail"]
        )

        # Give a quick indicator of whether this mail is an inbox mail or sent mail:
        all_recipients = []
        for field in ["to", "cc", "bcc"]: all_recipients += [item["email"] for item in client_response.data[field]]
        if google_tokens.email in all_recipients: mail_message.isSent = False
        else: mail_message.isSent = True

        # Invoke the LLM:
        ai_snippet = await self.summarize_mail_with_ai(
            mongo_conn = mongo_conn,
            user_info = user_info,
            llm = llm,
            message = mail_message
        )
        mail_message.aiSnippet = ai_snippet.summary

        # Done here:
        sync_result.success = True
        sync_result.mailMessage = mail_message
        return sync_result

    async def __sync_many_gmail(
            self,
            db_conn: AsyncMySQL,
            mongo_conn: AsyncMongo,
            user_info: CoreUserInfoModel,
            auth_token: CoreAuthTokenModel,
            mail_client: AsyncGMailClient,
            llm: CoreLLMController = None,
            force_sync: bool = False,
            start_date: datetime.datetime = None,
            end_date: datetime.datetime = None,
            max_count: int = 100,
            session_token: str = None
    ) -> MailSyncManyResults:

        # Start by assuming failure:
        sync_results = MailSyncManyResults()

        # Extract the client's tokens from the full token payload given by the database:
        google_tokens = GoogleAuthTokens(**auth_token.token)

        # Refresh the tokens (if needed):
        tokens_refreshed = await google_tokens.arefresh(
            http_client = current_app.http_client,
            client_id = mail_client.client_id,
            client_secret = mail_client.client_secret
        )
        if tokens_refreshed:
            auth_token.token = google_tokens.model_dump()
            auth_token.lastRefreshTs = date_time.get_current_utc_date_time(as_string = True)
            await self.set_token(
                db_conn = db_conn,
                mongo_conn = mongo_conn,
                token_key = auth_token.key,
                auth_token = auth_token,
                session_token = session_token
            )

        # Let's build the query to send to Google:
        sub_queries = []
        if start_date: sub_queries.append(start_date.strftime("after:%Y/%m/%d"))
        if end_date: sub_queries.append((end_date + datetime.timedelta(days = 1)).strftime("before:%Y/%m/%d"))
        query_string = " ".join(sub_queries)

        # Let's enlist all the mails that fall in the date range:
        client_response = await mail_client.list_messages(
            tokens = google_tokens,
            max_count = max_count,
            query = query_string
        )
        if not client_response.success:
            sync_results.message = f"gmail: {client_response.message}"
            return sync_results
        messages_list = client_response.data["messages"]

        # Now, for every mail in the list, we fetch the mail and note the results:
        tasks = [
            self.__sync_one_gmail(
                mongo_conn = mongo_conn,
                user_info = user_info,
                auth_token = auth_token,
                mail_client = mail_client,
                google_tokens = google_tokens,
                message_id = v["id"],
                llm = llm,
                force_sync = force_sync
            ) for v in messages_list.values()
        ]
        individual_sync_results = await asyncio.gather(*tasks)

        # Now we create operations for each mail,
        # and maintain success/failure counters:
        sync_results.totalCount = len(individual_sync_results)
        mongo_operations = []
        for result in individual_sync_results:
            if result.success: sync_results.successCount += 1
            else: sync_results.failureCount += 1
            if result.mailMessage:
                replacement_json = result.mailMessage.model_dump()
                replacement_json.pop("_id", None)
                mongo_operations.append(
                    ReplaceOne(
                        filter = {
                            "tokenId": ObjectId(auth_token.authTokenId),
                            "serviceType": auth_token.serviceType,
                            "client": auth_token.client,
                            "clientMessageId": result.mailMessage.clientMessageId
                        },
                        replacement = replacement_json,
                        upsert = True
                    )
                )

        # Make the bulk insert operation:
        if mongo_operations:
            sync_count = await current_app.core_message_controller.bulk_write(
                mongo_conn = mongo_conn,
                mongo_operations = mongo_operations
            )

        # Apply the labels to the read messages:
        try:
            client_response = await mail_client.modify_messages(
                tokens = google_tokens,
                message_ids = [v["id"] for v in messages_list.values()],
                add_label_ids = [google_tokens.labels.get("TCAOFF", {}).get("id")]
            )
        except Exception as exception:
            pass

        # Done here:
        sync_results.message = f"{sync_results.successCount}/{sync_results.totalCount} mail(s) sync'd from gmail"
        return sync_results

    async def sync(
            self,
            db_conn: AsyncMySQL,
            mongo_conn: AsyncMongo,
            user_info: CoreUserInfoModel,
            token_key: ObjectId | str,
            llm: CoreLLMController = None,
            force_sync: bool = False,
            start_date: datetime.datetime = None,
            end_date: datetime.datetime = None,
            max_count: int = 100,
            session_token: str = None
    ) -> MailSyncManyResults:

        # Start by assuming failure:
        sync_results = MailSyncManyResults()

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

        # We first load the authorization tokens:
        auth_token = await self.get_token(
            mongo_conn = mongo_conn,
            token_key = token_key,
        )

        # If we failed to load the authorization tokens:
        if not auth_token:
            sync_results.message = f"no such token key '{token_key}'"
            return sync_results

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

        if auth_token.client == "gmail":
            return await self.__sync_many_gmail(
                db_conn = db_conn,
                mongo_conn = mongo_conn,
                user_info = user_info,
                auth_token = auth_token,
                mail_client = current_app.gmail_client,
                llm = llm,
                force_sync = force_sync,
                start_date = start_date,
                end_date = end_date,
                max_count = max_count,
                session_token = session_token,
            )

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

        # If we haven't been able to sync mail due to not entering any 'if' condition:
        sync_results.message = f"no such mail client '{auth_token.client}'"
        return sync_results

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

    # These are simply for retrieving mails. You need to already have them sync'd to the database. These methods don't
    # fetch the mails from the third-party clients.

    @staticmethod
    async def list_mails(
            mongo_conn: AsyncMongo,
            token_ids: List[ObjectId | str],
            limit: int = 100,
            skip: int = 0,
            additional_filter: dict = None
    ) -> List[CoreMessageModel] | None:

        # regardless of what additional filter is provided from outside,
        # we add a mail-selecting filter here:
        if additional_filter is None: additional_filter = {}
        additional_filter["serviceType"] = "email"

        # Simply call the core model:
        return await current_app.core_message_controller.get_previews(
            mongo_conn = mongo_conn,
            token_ids = token_ids,
            limit = limit,
            skip = skip,
            additional_filter = additional_filter
        )

    @staticmethod
    async def get_one_mail(
            mongo_conn: AsyncMongo,
            token_id: ObjectId | str,
            message_id: ObjectId | str
    ) -> CoreMessageModel | None:

        # Simply call the core model:
        return await current_app.core_message_controller.get_message(
            mongo_conn = mongo_conn,
            token_id = token_id,
            message_id = message_id
        )

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

    @staticmethod
    async def update_tags(
            mongo_conn: AsyncMongo,
            token_id: ObjectId | str,
            message_id: ObjectId | str,
            unset_tags: List[str] = None,
            set_tags: List[str] = None
    ) -> bool:

        # Simply call the core model:
        return await current_app.core_message_controller.update_tags(
            mongo_conn = mongo_conn,
            token_id = token_id,
            message_id = message_id,
            unset_tags = unset_tags,
            set_tags = set_tags
        )


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


if __name__ == "__main__":

    pass

    # from utils_v2.string import json
    #
    # file_options = [
    #     r"/home/developer/Downloads/recursive parts parse - 20241210.json",
    #     r"/home/developer/Downloads/recursive parts parse (no attachment) - 20241210.json",
    # ]
    #
    # raw_mail_json = json.from_file(file_options[1])
    # print("FROM FILE:", json.to_string(raw_mail_json["payload"]))
    # print("\n\n---------\n\n")
    # mail_controller = MailController()
    # print(json.to_string(mail_controller.drop_attachments(raw_mail_json["payload"])))
