# -*- coding: utf-8 -*-
# @Author: Vivian Li
# @Date:   2024-02-24 10:13:41
# @Email:
# @Last Modified by:   Vivian Li
# @Last Modified time: 2024-04-11 00:23:34
# @Descriptions : A source will be converted to a payload object at the beginning
#                 of the transcription process.
#                 Given a list of datafiles, a payload object will implement
#                 function to transcribe and analyse the sources with the specified profile setting


import os.path
from abc import ABC, abstractmethod
from copy import deepcopy
from dataclasses import dataclass
from datetime import datetime
from typing import List

from gailbot.pluginSuiteManager.suite.gbPluginMethod import GBPluginMethods
from gailbot.profileManager import ProfileObject
from gailbot.shared.utils.general import make_dir, get_name, write_json, is_file, copy
from gailbot.shared.utils.logger import makelogger
from gailbot.sourceManager.watcher import Watcher
from gailbot.payload.result import SttTranscribeResult, AnalysisResult, FormatResult
from gailbot.workspace.directory_structure import OutputFolder, TemporaryFolder
import time

logger = makelogger("payloadObject")


@dataclass
class Progress:
    Start = "\u25B6 {name} : Start"
    Waiting = "\U0001F551  {name} : Waiting"
    Transcribing = "\U0001F50A {name} : Transcribing with engine {engine}"
    Finished = "\u2705 {name}: Completed"
    Transcribed = "\u2705 {name}: Transcribed"
    Error = "\U0001F6AB {name}: Error"
    Analyzing = "\U0001F4AC {name}: Applying plugin suite {suite}"
    Analyzed = "\u2705 {name}: Analyzed"
    Formatting = "\U0001F4C1 {name}: Formatting"


class PayloadObject(ABC):
    def __init__(
        self,
        original_source: str,
        data_files: List[str],
        temp_dir_struct: TemporaryFolder,
        output_path: str,
        profile: ProfileObject,
        watchers: List[Watcher],
    ):
        """
        Parameters
        ----------
        original_source:  The original path to the initial source object from user, source will be written to meta.json
        data_files: a list of datafiles, different payload may accept different datafiles format,
                    if it is a path to a media file, the file format will be in wav
        temp_dir_struct: stores paths to temporary directory, all paths in temp_dir should be already a directory
        output_path: path to output directory
        profile: Profile setting applied to the transcription
        watchers: A list of watcher added by user for the source to monitor the transcription progress
        """
        self.data_files = data_files
        self.original_source = original_source
        self.name = get_name(original_source)
        self.profile = profile
        # payload_id is generated by the time stamp
        self.payload_id = self.name + "_" + datetime.now().strftime("%m%d_%H_%M_%S_%f")

        self.output_dir = os.path.join(output_path, self.payload_id)
        self.temp_dir_struct = temp_dir_struct
        self.watchers = deepcopy(watchers)
        self.transcription_result: SttTranscribeResult = SttTranscribeResult(
            self.temp_dir_struct.transcribe_ws
        )
        self.analysis_result: AnalysisResult = AnalysisResult(
            self.temp_dir_struct.analysis_ws
        )
        self.format_result: FormatResult = FormatResult(self.temp_dir_struct.format_ws)
        self.profile = profile
        self.engine = profile.engine_provider.make_engine()
        self.plugin_suites = profile.plugin_suites


        self.source_dir = self.temp_dir_struct.data_copy

        # merged_audio and output_dir_struct will not be initialized until transcription starts
        self.merged_audio: str | None = None
        self.output_dir_struct: OutputFolder | None = None
        self._starting_time = 0

    @abstractmethod
    def stt_transcribe(self) -> bool:
        """
        subclass is responsible to implement this function that after it is invoked, the result of transcribing
        [self.datafiles] will be stored in [self.transcription_result]

        Returns
        -------
        return true of the transcription succeed, false otherwise
        """
        raise NotImplementedError()

    @abstractmethod
    def merge_audio(self, output) -> str:
        """
        subclass is responsible to implement this function that based on the structure of [self.data_files],
        merge all audio files in the payload to a single file stored in the [output] directory
        Parameters
        ----------
        output

        Returns
        -------
        path where the merged audio file is stored
        """
        raise NotImplementedError()

    def send_msg_to_watchers(self, content: str):
        for w in self.watchers:
            w.watch(progress=f"{self.name} : {content}")

    def analyze(self) -> bool:
        """
        perform transcription analysis using the applied plugin suite

        Returns
        -------
        return true if the analysis succeed, false otherwise

        Notes
        -------
        this function will only be invoked when stt_transcribe return true,
        """
        try:
            results = dict()
            for suite, selected in self.plugin_suites.items():

                suite_output = os.path.join(self.output_dir_struct.analysis, suite.name)
                suite_ws = os.path.join(self.temp_dir_struct.analysis_ws, suite.name)
                make_dir(suite_output)
                make_dir(suite_ws)
                for watcher in self.watchers:
                    watcher.watch(
                        Progress.Analyzing.format(name=self.name, suite=suite.name)
                    )
                result = suite(
                    base_input=None,
                    methods=GBPluginMethods(
                        work_path=suite_ws,
                        out_path=suite_output,
                        data_files=self.data_files,
                        merged_media=self.merged_audio,
                        utterances=self.transcription_result.get_data(),
                    ),
                    selected_plugins=selected,
                ) #nothing is being returned by suite.__call__ because of summer changes. 
                results[suite.name] = result
            assert self.analysis_result.save_data(results)
            return True
        except Exception as e:
            logger.error(e, exc_info=True)
            return False

    def format(self) -> bool:
        """
        Output the transcription result

        Returns
        -------

        Notes
        -------
        this function will only be invoked when both analysis and stt_transcribe return true,
        """
        try:
            assert self.analysis_result.output(self.output_dir_struct.analysis)
            assert self.transcription_result.output(self.output_dir_struct.transcribe)
            assert self.format_result.output(self.output_dir_struct.about)
            metadata = [
                {
                    "Profile Setting": {
                        "engine_name": self.profile.profile_data.engine_setting_name,
                        "suite_setting": self.profile.profile_data.plugin_suite_setting,
                    },
                    "Date": datetime.now().strftime("%Y-%m-%d %H:%M"),
                    "Source": self.original_source,
                    "Plugin Results": self.analysis_result.get_serialized_data(),
                    "time_elapsed": f"{time.process_time() - self._starting_time} seconds",
                }
            ]

            metafile = os.path.join(self.output_dir_struct.about, "meta.json")
            write_json(metafile, metadata)

            with open(
                os.path.join(
                    self.output_dir_struct.root, OutputFolder.GB_RESULT_SIGNATURE
                ),
                "w+",
            ) as f:
                f.write(f"{self.name}")
            for file in self.data_files:
                if is_file(file):
                    copy(
                        file,
                        os.path.join(
                            self.output_dir_struct.media, os.path.basename(file)
                        ),
                    )
            for suite, selected in self.plugin_suites.items():
                analysis = os.path.join(self.temp_dir_struct.analysis_ws, suite.name, "final")
                print("HIU", is_file(analysis))
                print(copy(analysis, os.path.join(self.output_dir_struct.analysis, suite.name, "final_analysis.xml"),))
            return True
        except Exception as e:
            logger.error(e, exc_info=True)
            return False

    def get_transcription_result(self):
        return self.transcription_result.get_data()

    def start_execute(self) -> bool:
        logger.info(f"Start transcribing {self.name}")
        self._starting_time = time.process_time()
        # initializing output directory and create merged audio file
        self.output_dir_struct = OutputFolder(self.output_dir)
        # since merged_audio can be used by plugin suite, it needs to be initialized before analyze phase
        self.merged_audio = self.merge_audio(self.output_dir_struct.media)
        logger.info(f"reached merged_audio")

        self.output_dir = self.output_dir_struct.root
        transcribed = self.stt_transcribe()
        logger.info(f"passed stt_transcribe")

        if transcribed:
            analyzed = self.analyze()
            print("\n\n\nAnalyzed! Results: ", analyzed)
        else:
            logger.error(f"{self.name} failed transcription")
            return False

        # if analyzed:
        for watcher in self.watchers:
            watcher.watch(Progress.Formatting.format(name=self.name))
        formatted = self.format()
        # else:
        #     logger.error(f"{self.name} failed analysis")
        #     return False

        if formatted:
            for watcher in self.watchers:
                watcher.watch(Progress.Finished.format(name=self.name))
            return True
        else:
            logger.error(f"{self.name} failed formatting")
            return False
