import os
import io
import sys
import time
import gradio as gr
from typing import List, Union, Any, Dict, Optional
import torch
import numpy as np
import librosa
from transformers import SpeechSynthesisPipeline, AutoTokenizer, Wav2Vec2ForCTC, Wav2Vec2Processor
from pypdfium2 import PdfDocument, Page

def convert_pdf_to_text(filepath):
    doc = PdfDocument(filepath)
    text = ""
    for page in doc.pages():
        text += page.extract_text() + "\n\n"
    return text

class QuantizedSpeechT5TTSPipe:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained('facebook/wav2vec2-base-960h')
        self.model = Wav2Vec2ForCTC.from_pretrained('/model/quantized_model').half().cuda()
        self.processor = Wav2Vec2Processor.from_pretrained("/model/quantized_vocab")
        
    def _pad_and_concatenate(self, tensor_list: List[torch.Tensor], padding_value=0):
        max_length = max([tensor.size(0) for tensor in tensor_list])
        padded_tensors = []
        for tensor in tensor_list:
            pad_width = max_length - tensor.size(0)
            if pad_width > 0:
                tensor_padded = torch.cat((tensor, torch.full((pad_width, tensor.size(1)), fill_value=padding_value).type_as(tensor)))
            else:
                tensor_padded = tensor
            padded_tensors.append(tensor_padded)
        
        return torch.stack(padded_tensors)
    
    def preprocess(self, inputs: Union[str, List[str]], **kwargs) -> dict:
        if isinstance(inputs, str):
            inputs = [inputs]
        batch_encodings = self.tokenizer(inputs, truncation=True, padding='longest', return_tensors="pt").input_values
        return {"batch_encodings": batch_encodings}
    
    def postprocess(self, outputs: Dict[str, torch.Tensor], **kwargs) -> Union[List[str], List[bytes]]:
        logits = outputs["logits"].cpu().detach().numpy()
        ids = np.argmax(logits, axis=-1)
        cleaned_ids = [id_seq[:np.where(np.array(id_seq) == 2)[0][0]] for id_seq in ids] # Remove CTC blanks
        decoded_strings = self.tokenizer.decode(cleaned_ids)
        
        audios = []
        for text in decoded_strings:
            input_values = self.processor(text, sampling_rate=16000, return_tensors="pt").input_values
            input_values = input_values.cuda().unsqueeze(0)
            
            mel_outputs = self.model(input_values).mel_output
            _, predicted_ids = torch.topk(mel_outputs.float(), k=1, dim=-1)
            predicted_ids = predicted_ids.squeeze(-1).tolist()[0]
            
            raw_waveform = self.processor.post_processing(predicted_ids)
            waveform = raw_waveform * 32768 / max(abs(raw_waveform))
            wav_data = np.int16(waveform)
            audio = io.BytesIO()
            sf.write(audio, int(44100), wav_data, format="WAV")
            audios.append(audio.getvalue())
        return audios
    
    def generate(self, text: str):
        processed_inputs = self.preprocess(text)
        outputs = self.model(**processed_inputs)
        results = self.postprocess(outputs)
        return results

if __name__ == "__main__":
    tts = QuantizedSpeechT5TTSPipe()
    sample_text = 'Hello world! This is a test.'
    result = tts.generate(sample_text)
    print(f'Generated {len(result)} audio files from "{sample_text}"')

def main(pdf_file: gr.File, output_filename: str):
    start_time = time.time()
    pdf_text = convert_pdf_to_text(pdf_file)
    print(f'Processed PDF content in {time.time() - start_time:.4f} seconds')

    pipe = QuantizedSpeechT5TTSPipe()
    start_time = time.time()
    audios = pipe.generate(pdf_text)
    print(f'Generated {len(audios)} audio files in {time.time() - start_time:.4f} seconds')

    zip_buffer = BytesIO()
    with ZipFile(zip_buffer, mode='w') as zf:
        for i, audio in enumerate(audios):
            filename = f"{i}_{output_filename}.wav"
            zf.writestr(filename, audio)
    zip_buffer.seek(0)

    return {'zip': zip_buffer}

iface = gr.Interface(fn=main, 
                      inputs="file", 
                      outputs="binary",
                      input_types=['pdf'],
                      output_types=['download'])

iface.launch()