"""
stg - Scraibe to Gradio Interface

This module provides an interface between the Scraibe transcription system and the Gradio user interface.
It defines a class, ScraibeWrapper, that wraps the Scraibe model and provides methods for performing transcription tasks through the Gradio UI.

Modules:
    json: Used for encoding and decoding JSON data.
    gradio as gr: Used for creating the Gradio UI.
    tqdm: Used for displaying progress bars.
    scraibe.app.global_var as gv: Contains global variables for the Scraibe app.
"""
import json
import gradio as gr
from tqdm import tqdm
from typing import Any, Dict, Union, Tuple, List

from scraibe import Scraibe, Transcriber

class ScraibeWrapper:
    """
    A class that provides an interface between the Gradio UI and the Scraibe transcription system.

    This class wraps the Scraibe model and provides methods for performing transcription tasks through the Gradio UI. 
    These tasks include auto transcription, transcription, and diarisation.

    Attributes:
        model (Scraibe): The Scraibe model for performing transcription tasks.
    """

    def __init__(self, model) -> None:
        """
        Initializes the ScraibeWrapper with a Scraibe model.

        Args:
            model (Scraibe): The Scraibe model for performing transcription tasks.
            *args (Any): Additional positional arguments.
            **kwargs (Dict[str, Any]): Additional keyword arguments.
        """
            
        self.model = model

    def autotranscribe(self, source: Union[str, List[str]],
                        num_speakers: int,
                        translate: bool,
                        language: str,
                        *args: Any, **kwargs: Dict[str, Any]) -> Tuple[str, Union[str, dict]]:
        """
        Performs auto transcription on the given source.

        Args:
            source (Union[str, List[str]]): The source to transcribe. This can be a string representing a single source,
                                    or a list of strings representing multiple sources.
            num_speakers (int): The number of speakers in the source.
            translate (bool): Whether to translate the transcription.
            language (str): The language of the source.
            *args (Any): Additional positional arguments.
            **kwargs (Dict[str, Any]): Additional keyword arguments.

        Returns:
            Tuple[str, Union[str, dict]]: A tuple containing the transcribed text (str) and the JSON output (str or dict).
        """
        
        _kwargs = {
            "num_speakers": num_speakers if num_speakers != 0 else None,
            "language": language if language != "None" else None,
            "task": 'translate' if translate else None
        }
        if isinstance(source, str):
            try:
                result = self.model.autotranscribe(source, **_kwargs)
            except ValueError:
                raise gr.Error("Couldn't detect any speech in the provided audio. \
                        Please try again!")
                
            return result, str(result), result.get_json()
        
        elif isinstance(source, list):
            source_names = [s.split("/")[-1] for s in source]
            result = []
            for s in tqdm(source, total=len(source),desc = "Transcribing audio files"):
                try:
                    res = self.model.autotranscribe(s, **_kwargs)
                except ValueError:
                    _name = s.split("/")[-1]
                    res = f"NO TRANSCRIPT FOUND FOR {_name}"
                    gr.Warning(f"Couldn't detect any speech in {_name} will skip this file.")
                result.append(res)
            
            out = ''
            out_dict = {}
            for i, r in enumerate(result):
                out += f"TRANSCRIPT FOR {source_names[i]}:\n\n"
                out += str(r)
                out += "\n\n"
                
                if isinstance(r, str):
                    out_dict[source_names[i]] = r
                else:
                    out_dict[source_names[i]] = r.get_dict()
              
            return out, json.dumps(out_dict, indent=4)
        
        else:
            raise gr.Error("Please provide a valid audio file.")


    def transcribe(self, source: Union[str, List[str]], 
                   translate: bool,
                   language: str,
                   *args: Any, **kwargs: Dict[str, Any]) -> str:
        """
        Performs transcription on the given source.

        Args:
            source (Union[str, List[str]]): The source to transcribe.
            This can be a string representing a single source, or a list of strings representing multiple sources.
            translate (bool): Whether to translate the transcription.
            language (str): The language of the source.
            *args (Any): Additional positional arguments.
            **kwargs (Dict[str, Any]): Additional keyword arguments.

        Returns:
            str: The transcribed text.
        """

        _kwargs = {
            "language": language if language != "None" else None,
            "task": 'translate' if translate == "Yes" else None
        }
    
        if isinstance(source, str):
            result = self.model.transcribe(source, **_kwargs)

            return str(result)
        
        elif isinstance(source, list):
            source_names = [s.split("/")[-1] for s in source]
            result = []
            for s in tqdm(source, total=len(source),desc = "Transcribing audio files"):
                res = self.model.transcribe(s, **_kwargs)
                result.append(res)
            
            out = ''
            for i, res in enumerate(result):
                out += f"TRANSCRIPT FOR {source_names[i]}:\n\n"
                out += str(res)
                out += "\n\n"
            
            return out
        
        else:
            raise gr.Error("Please provide a valid audio file.")

    def diarisation(self, source: Union[str, List[str]],
                    num_speakers: int,
                    *args: Any, **kwargs: Dict[str, Any]) -> str:
        """
        Performs diarisation on the given source.

        Args:
            source (Union[str, List[str]]): The source to perform diarisation on.
            This can be a string representing a single source, 
            or a list of strings representing multiple sources.
            num_speakers (int): The number of speakers in the source.
            *args (Any): Additional positional arguments.
            **kwargs (Dict[str, Any]): Additional keyword arguments.

        Returns:
            str: The JSON output of the diarisation result.
        """
    
        
        _kwargs = {
            "num_speakers": num_speakers if num_speakers != 0 else None,
        }
        
        if isinstance(source, str):
            try:
                result = self.model.diarization(source, **_kwargs)
            except ValueError:
                raise gr.Error("Couldn't detect any speech in the provided audio. \
                        Please try again!")
                
            return json.dumps(result, indent=2)
        elif isinstance(source, list):
            source_names = [s.split("/")[-1] for s in source]
            result = []
            for s in tqdm(source, total=len(source),desc = "Performing diarisation"):
                try:
                    res = self.model.diarization(s, **_kwargs)
                except ValueError:
    
                    res = f"NO DIARISATION FOUND FOR {s}"
                    gr.Warning(f"Couldn't detect any speech in {s} will skip this file.")
                result.append(res)
            
            out = {}
            
            for i, res in enumerate(result):
                out[source_names[i]] = res
            
            return json.dumps(out, indent=4)
        
        else:
            gr.Error("Please provide a valid audio file.")
    
    def get_task_from_str(self, task: str) -> callable:
        """
        Returns the corresponding task function based on the given task string.

        Args:
            task (str): The task string. This can be one of the following: 'Auto Transcribe', 'Transcribe', 'Diarisation'.

        Returns:
            callable: The corresponding task function.
        """
        
        if task == 'Auto Transcribe':
            return self.autotranscribe
        elif task == 'Transcribe':
            return self.transcribe
        elif task == 'Diarisation':
            return self.diarisation
        else:
            raise ValueError("Invalid task string.")
    
    def update_transcriber_model(self, model: str, **kwargs) -> None:
        """
        Load the selected model into memory.

        Args:
            model (str): The name of the model to load.
            kwargs (Dict[str, Any]): Additional keyword arguments.
        """
        
        self.model.transcriber = Transcriber.load_model(model, **kwargs)
    
    @classmethod
    def load_from_dict(cls, config: Dict[str, Any]) -> 'ScraibeWrapper':
        """ Load the ScraibeWrapper from a dictionary configuration.
        
        Args:
            config (dict): A dictionary containing the configuration parameters.
            
        Returns:
            ScraibeWrapper: The ScraibeWrapper object.
        
        """
        model = Scraibe(**config)
        
        return cls(model)

        
        
               
