import gradio as gr
import fitz  # PyMuPDF
import edge_tts
import asyncio
import os
import uuid
from pydub import AudioSegment

# Fetch available voices
async def get_voices():
    voices = await edge_tts.VoicesManager.create()
    voice_list = voices.find(Language="en")
    return {f"{v['FriendlyName']} ({v['ShortName']})": v['ShortName'] for v in voice_list}

# Function to extract text from PDF
def extract_text(pdf_file):
    doc = fitz.open(pdf_file.name)
    return "".join([page.get_text() for page in doc])

# Main conversion logic with Progress Tracker
async def convert_pdf_to_long_audio(pdf_file, voice_short_name, progress=gr.Progress()):
    if pdf_file is None:
        return "Please upload a file.", None

    progress(0, desc="Reading PDF...")
    text = extract_text(pdf_file)
    
    if not text.strip():
        return "No text found in PDF.", None

    # Chunking: ~2500 characters is a safe bet for Edge-TTS stability
    chunk_size = 2500
    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
    total_chunks = len(chunks)
    
    combined_audio = AudioSegment.empty()
    session_id = str(uuid.uuid4())
    os.makedirs(session_id, exist_ok=True)

    # Iterating through chunks with progress updates
    for i, chunk in enumerate(chunks):
        # Update progress bar: (current_index / total_count)
        progress((i / total_chunks), desc=f"Converting chunk {i+1} of {total_chunks} to voice...")
        
        chunk_path = os.path.join(session_id, f"chunk_{i}.mp3")
        communicate = edge_tts.Communicate(chunk, voice_short_name)
        await communicate.save(chunk_path)
        
        # Load and append
        segment = AudioSegment.from_mp3(chunk_path)
        combined_audio += segment

    progress(0.95, desc="Merging all audio parts into final file...")
    final_path = f"audiobook_{session_id}.mp3"
    combined_audio.export(final_path, format="mp3")
    
    # Cleanup
    for f in os.listdir(session_id):
        os.remove(os.path.join(session_id, f))
    os.rmdir(session_id)
    
    progress(1.0, desc="Done!")
    return text[:2000] + "...", final_path

# Gradio Wrapper
def process(pdf, voice_name, progress=gr.Progress()):
    voice_id = voice_dict[voice_name]
    return asyncio.run(convert_pdf_to_long_audio(pdf, voice_id, progress))

# Initialize voice list
voice_dict = asyncio.run(get_voices())

# Building the Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎧 Infinite PDF Audiobook Generator")
    gr.Markdown("Upload your PDF and wait for the AI to narrate it. **Progress is tracked below the button.**")
    
    with gr.Row():
        with gr.Column():
            pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
            voice_input = gr.Dropdown(
                choices=list(voice_dict.keys()), 
                label="Select AI Voice", 
                value="Microsoft Guy Online (Natural) - en-US-GuyNeural"
            )
            btn = gr.Button("Start Audio Conversion", variant="primary")
        
        with gr.Column():
            text_preview = gr.Textbox(label="Text Preview", lines=5)
            audio_output = gr.Audio(label="Final Audiobook (Download here)")

    # The magic happens here: passing the progress bar to the function
    btn.click(process, inputs=[pdf_input, voice_input], outputs=[text_preview, audio_output])

if __name__ == "__main__":
    demo.launch()