import time
from typing import Any, Union
from rich.console import Console
from rich.table import Table

from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.Core.Utilities.ReturnValues import DOKReturnType, DErrorReturnType
from DIRAC.FrameworkSystem.Client.NotificationClient import NotificationClient
from DIRAC.TransformationSystem.Client.TransformationClient import (
    TransformationClient,
)
from DIRAC.TransformationSystem.Agent.TransformationAgentsUtilities import (
    TransformationAgentsUtilities,
)
from DIRAC.ProductionSystem.Client.ProductionClient import ProductionClient

from CTADIRAC.Core.Utilities.return_values import DReportReturnType, s_report


class FailoverUtilities(TransformationAgentsUtilities):
    """
    Some useful methods used by TransformationFailoverAgent
    """

    def __init__(self, agent_name) -> None:
        TransformationAgentsUtilities.__init__(self)
        self.notif_client = NotificationClient()
        self.trans_client = TransformationClient()
        self.prod_client = ProductionClient()

        self.mail_from = None
        self.mail_to = None
        self.enabled = False

    def get_eligible_transformation(
        self, status, type_list
    ) -> Union[Any, DOKReturnType]:
        """Select transformations of given status and type.
        using the TransformationClient"""
        res = self.trans_client.getTransformations(
            condDict={"Status": status, "Type": type_list}
        )
        if not res["OK"]:
            return res
        transformations = {}
        for trans_info in res["Value"]:
            trans_id = trans_info["TransformationID"]
            transformations[str(trans_id)] = trans_info
        return S_OK(transformations)

    def get_transformation_tasks(self, trans_id: int) -> Union[Any, DOKReturnType]:
        """get transformations tasks
        using the TransformationClient"""
        res = self.trans_client.getTransformationTasks({"TransformationID": trans_id})
        if not res["OK"]:
            gLogger.error(res["Message"])

        if len(res["Value"]) == 0:
            gLogger.notice(f"[{trans_id}] No tasks selected")

        return res

    def get_transformation_files(self, cond_dict) -> Union[Any, list]:
        """Get the transformation files
        using the TransformationClient"""

        res = self.trans_client.getTransformationFiles(cond_dict)
        if res["OK"]:
            return res["Value"]
        else:
            gLogger.error(res)
            return res["Value"]

    def is_transformation_processed(self, trans_id: int) -> bool:
        res_trans_stat = self.trans_client.getTransformationStats(trans_id)
        is_processed = False
        if res_trans_stat["OK"]:
            try:
                if (
                    res_trans_stat["Value"]["Processed"]
                    == res_trans_stat["Value"]["Total"]
                ):
                    is_processed = True
            except KeyError:
                gLogger.info(f"[{trans_id}] Transformation got no 'Processed' files.")
        else:
            gLogger.error(
                f"[{trans_id}] Can't getTransformationStats: {res_trans_stat}"
            )
        return is_processed

    def complete_transformation(
        self, trans_id: int
    ) -> Union[DReportReturnType, DErrorReturnType]:
        res_complete = self.trans_client.completeTransformation(trans_id)
        if res_complete["OK"]:
            return s_report(False, f"[{trans_id}] Transformation Completed.")
        else:
            return S_ERROR(
                f"[{trans_id}]: Can't complete transformation: {res_complete}"
            )

    def get_parent_transformation_id(self, trans_id: int) -> Union[int, None]:
        res = self.prod_client.get_parent_transformation(trans_id)
        if res["OK"]:
            if res["Value"]:
                try:
                    parent_id = res["Value"][0][0]
                    if isinstance(parent_id, int):
                        return parent_id
                except (TypeError, IndexError):
                    gLogger.error(
                        f"[{trans_id}] Error while getting parent transformation: {res}"
                    )
                    return None
            else:
                gLogger.info(f"[{trans_id}] Parent ID has empty value")
                return None
        else:
            gLogger.error(
                f"[{trans_id}] Error while getting parent transformation: {res}"
            )
            return None

    def get_parent_type(self, parent_id: int) -> Union[Any, None]:
        res = self.trans_client.getTransformations(
            condDict={"TransformationID": parent_id}
        )
        if res["OK"]:
            return res["Value"][0]["Type"]

    def sorting_key(self, index1, index2):
        return lambda item: (item[index1], item[index2])

    def create_mail_message(
        self, transformation_reports: dict[str, dict[str, str]]
    ) -> str:
        """Sort the transformation by types
        and create the html string message"""

        message = ""
        for trans_type in transformation_reports:
            message += """<!DOCTYPE html>
            <html>
            <head>
            <meta charset="UTF-8">
            <style>
            .r1 {font-style: italic}
            .r2 {font-weight: bold}
            .r3 {color: #008000; text-decoration-color: #008000; font-weight: bold}
            .r4 {color: #800000; text-decoration-color: #800000; font-weight: bold}
            body {
                color: #000000;
                background-color: #ffffff;
            }
            </style>
            </head>
            <body>
                <pre style="font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"><code><span class="r1">"""
            message += f"Report on {trans_type} transformations:"
            message += """</span>
            </code></pre>
            </body>
            </html>"""
            for report in transformation_reports[trans_type].values():
                message += report + "\n"
        return message

    def init_rich_table(self, title, fields) -> Table:
        """Init Rich Table:"""
        table = Table(title=title)
        for f in fields:
            table.add_column(f)
        return table

    def generate_html_table(self, table) -> str:
        console = Console(record=True)
        console.print(table)
        html_output: str = console.export_html()
        return html_output

    def send_mail(self, subject, message, html=True) -> None:
        gLogger.info(f"Sending mail from {self.mail_from} to {self.mail_to}")
        res = self.notif_client.sendMail(
            self.mail_to, subject, message, self.mail_from, localAttempt=True, html=html
        )
        if res["OK"]:
            gLogger.info(f"Report succesfully sent to {self.mail_to}")
        else:
            gLogger.error(res)

    @classmethod
    def measure_time(cls, func):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            elapsed_time = end_time - start_time
            gLogger.info(f"{func.__name__} took {elapsed_time:.2f} seconds to execute")
            return result

        return wrapper
