import gradio as gr
from TTS.api import TTS
import fitz  # PyMuPDF for PDF text extraction
from tempfile import NamedTemporaryFile
import os

# Initialize TTS models
TTS_MODELS = {
    "American - Male": "tts_models/en/ljspeech/tacotron2-DDC",
    "American - Female": "tts_models/en/ljspeech/tacotron2-DDC",
}

def extract_text_from_pdf(pdf_path, page_selection=None):
    """Extract specific pages from a PDF."""
    pdf_document = fitz.open(pdf_path)
    text = ""
    total_pages = len(pdf_document)

    # Parse the page_selection input (e.g., "1,3-5")
    pages_to_read = []
    if page_selection:
        for part in page_selection.split(","):
            if "-" in part:
                start, end = map(int, part.split("-"))
                pages_to_read.extend(range(start - 1, end))  # Convert to 0-indexed
            else:
                pages_to_read.append(int(part) - 1)  # Convert to 0-indexed
    else:
        pages_to_read = list(range(total_pages))  # Default: all pages

    # Ensure valid page selection
    pages_to_read = [p for p in pages_to_read if 0 <= p < total_pages]

    # Extract text from selected pages
    for page_num in pages_to_read:
        page = pdf_document[page_num]
        page_text = page.get_text()
        text += f"\n--- Page {page_num + 1} ---\n{page_text}"  # Add page info

    if not text.strip():
        return "Error: No text found on the selected pages. They might contain images only."
    return text

def generate_audio(text, accent_gender, speed):
    """Generate audio from text using selected accent, gender, and speed."""
    model_name = TTS_MODELS[accent_gender]
    tts = TTS(model_name=model_name)

    # Generate audio and save as WAV
    with NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
        tts.tts_to_file(text=text, file_path=temp_audio.name)

        # Save as MP3
        output_mp3_path = temp_audio.name.replace(".wav", ".mp3")
        os.rename(temp_audio.name, output_mp3_path)  # Rename WAV to MP3
        return output_mp3_path

def process_input(input_text, uploaded_file, page_selection, accent_gender, speed):
    """Process input (text or file) and generate audio."""
    try:
        if not input_text and not uploaded_file:
            return "Please provide input text or upload a file.", None

        # Extract text from uploaded file
        if uploaded_file:
            file_extension = uploaded_file.name.split('.')[-1].lower()
            with NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(uploaded_file.read())
                temp_file_path = temp_file.name

            if file_extension == "pdf":
                text = extract_text_from_pdf(temp_file_path, page_selection)
            else:
                return "Error: Only PDF file support for page selection.", None

            os.remove(temp_file_path)
        else:
            text = input_text

        # Generate audio
        mp3_path = generate_audio(text, accent_gender, float(speed))
        return "Audio generated successfully!", mp3_path

    except Exception as e:
        import traceback
        error_message = f"Error: {str(e)}\n{traceback.format_exc()}"
        print(error_message)
        return error_message, None

# Gradio interface
interface = gr.Interface(
    fn=process_input,
    inputs=[
        gr.Textbox(label="Enter Text", placeholder="Type or paste text here...", lines=5),
        gr.File(label="Upload File (.pdf only for page selection)", file_types=[".pdf"]),
        gr.Textbox(label="Page Selection (e.g., '1,3-5' or leave blank for all pages)", placeholder="Pages to read"),
        gr.Dropdown(label="Accent & Gender", choices=list(TTS_MODELS.keys()), value="American - Male"),
        gr.Slider(label="Speed (e.g., 1.0 = Normal, 0.75 = Slower, 1.25 = Faster)", minimum=0.5, maximum=2.0, value=1.0, step=0.1),
    ],
    outputs=[
        gr.Textbox(label="Result"),
        gr.Audio(label="Generated Audio"),
    ],
    title="Text-to-Speech (TTS) Application",
    description="Upload a PDF file or enter text directly. Specify pages to extract text from, customize accent, gender, and speed. Download the generated audio as MP3."
)

# Launch the app
interface.launch()