from smolagents import Tool
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, logging
import warnings


class SpeechRecognitionTool(Tool):
    name = "speech_to_text"
    description = """Transcribes speech from audio."""

    inputs = {
        "audio": {
            "type": "string",
            "description": "Path to the audio file to transcribe.",
        },
        "with_time_markers": {
            "type": "boolean",
            "description": "Whether to include timestamps in the transcription output. Each timestamp appears on its own line in the format [float, float], indicating the number of seconds elapsed from the start of the audio.",
            "nullable": True,
            "default": False,
        },
    }
    output_type = "string"

    chunk_length_s = 30

    def __new__(cls, *args, **kwargs):
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

        model_id = "openai/whisper-large-v3-turbo"
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id,
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        )
        model.to(device)
        processor = AutoProcessor.from_pretrained(model_id)

        logging.set_verbosity_error()
        warnings.filterwarnings(
            "ignore",
            category=FutureWarning,
            message=r".*The input name `inputs` is deprecated.*",
        )
        cls.pipe = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            torch_dtype=torch_dtype,
            device=device,
            chunk_length_s=cls.chunk_length_s,
            return_timestamps=True,
        )

        return super().__new__(cls, *args, **kwargs)

    def forward(self, audio: str, with_time_markers: bool = False) -> str:
        """
        Transcribes speech from audio.

        Args:
            audio (str): Path to the audio file to transcribe.
            with_time_markers (bool): Whether to include timestamps in the transcription output. Each timestamp appears on its own line in the format [float], indicating the number of seconds elapsed from the start of the audio.

        Returns:
            str: The transcribed text.
        """
        result = self.pipe(audio)
        if not with_time_markers:
            return result["text"].strip()

        txt = ""
        for chunk in self._normalize_chunks(result["chunks"]):
            txt += f"[{chunk['start']:.2f}]\n{chunk['text']}\n[{chunk['end']:.2f}]\n"
        return txt.strip()

    def transcribe(self, audio, **kwargs):
        result = self.pipe(audio, **kwargs)
        return self._normalize_chunks(result["chunks"])

    def _normalize_chunks(self, chunks):
        chunk_length_s = self.chunk_length_s
        absolute_offset = 0.0
        chunk_offset = 0.0
        normalized = []

        for chunk in chunks:
            timestamp_start = chunk["timestamp"][0]
            timestamp_end = chunk["timestamp"][1]
            if timestamp_start < chunk_offset:
                absolute_offset += chunk_length_s
                chunk_offset = timestamp_start
            absolute_start = absolute_offset + timestamp_start

            if timestamp_end < timestamp_start:
                absolute_offset += chunk_length_s
            absolute_end = absolute_offset + timestamp_end
            chunk_offset = timestamp_end

            chunk_text = chunk["text"].strip()
            if chunk_text:
                normalized.append(
                    {
                        "start": absolute_start,
                        "end": absolute_end,
                        "text": chunk_text,
                    }
                )

        return normalized