import torch
from transformers import pipeline, AutomaticSpeechRecognitionPipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
from typing import List, Tuple, Any, Dict
from pydantic import BaseModel, Field



class Chunk(BaseModel):
    """
    Represents a chunk of text with associated timestamps.
    """
    timestamp: Tuple[float, float] = Field(..., description="The timestamp range for the chunk of text.")
    text: str = Field(..., description="The text content of the chunk.")

class Transcription(BaseModel):
    """
    Represents the transcription output, consisting of text and chunks.
    """
    text: str = Field(..., description="The entire text transcription.")
    chunks: List[Chunk] = Field(..., description="List of individual text chunks with timestamps.")

    @property
    def timestamps(self) -> List[Tuple[float, float]]:
        """
        Get the list of timestamps for each chunk.

        Returns:
            List[Tuple[float, float]]: List of timestamps for each chunk.
        """
        return [chunk.timestamp for chunk in self.chunks]

    @property
    def chunk_texts(self) -> List[str]:
        """
        Get the list of texts for each chunk.

        Returns:
            List[str]: List of texts for each chunk.
        """
        return [chunk.text for chunk in self.chunks]


class SpeechInference:
    """
    Class for transcribing speech using the Hugging Face Transformers library.
    """

    def __init__(self, model_name: str, huggingface_read_token: str) -> None:
        """
        Initialize the SpeechTranscriber instance.

        Args:
            model_name (str): The name of the Hugging Face model to use for transcription.
            huggingface_read_token (str): The Hugging Face API token for model access.
        """
        self.model_name = model_name
        self.huggingface_read_token = huggingface_read_token
    
    def pipe_initialization(self) -> AutomaticSpeechRecognitionPipeline:
        """
        Initialize the pipeline for speech transcription.

        Returns:
            Any: The initialized pipeline object.
        """
        if input is None:
            print("No audio file submitted! Please upload or record an audio file before submitting your request.")
        else:
            device = 0 if torch.cuda.is_available() else "cpu"
            torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
        self.model_name , torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
        )
        model.to(device)
        processor = AutoProcessor.from_pretrained(self.model_name)
        pipe = pipeline(
                        task="automatic-speech-recognition",
                        model=model,
                        tokenizer=processor.tokenizer,
                        feature_extractor=processor.feature_extractor,
                        token=self.huggingface_read_token,
                        device=device,
                        torch_dtype=torch_dtype
                        )
        return pipe
    
    def output(self, pipe: AutomaticSpeechRecognitionPipeline, input: Any, task: str) -> Dict:
        """
        Perform speech transcription.

        Args:
            pipe (Any): The initialized pipeline object.
            input (Any): The input data to transcribe.
            task (str): The speech task, e.g., "transcribe" or "translate".

        Returns:
            Transcription: The transcription output.
        """
        transcription = pipe(
                input,
                chunk_length_s=30,
                batch_size=24,
                return_timestamps=True,
                generate_kwargs={"task": task}
            )
        # transcription = Transcription(**transcription)
        return transcription





