import gradio as gr
import torch
import json
import numpy as np
import os
from datetime import datetime
from model import Image2Phoneme
from utils import ctc_post_process, audio_to_mel, mel_to_image, text_to_phonemes
import soundfile as sf
import shutil
import time

# Configuration
DEVICE = torch.device("cpu")
PHMAP = "phoneme_to_id.json"
AUDIO_DIR = "audio_inputs"

# Ensure audio directory exists
os.makedirs(AUDIO_DIR, exist_ok=True)

# Load phoneme vocabulary
try:
    vocab = json.load(open(PHMAP, "r"))
    id_to_ph = {v: k for k, v in vocab.items()}
except FileNotFoundError:
    raise FileNotFoundError(f"Phoneme mapping file not found at {PHMAP}")

# Build model
vocab_size = max(vocab.values()) + 1
model = Image2Phoneme(vocab_size=vocab_size).to(DEVICE)
try:
    ckpt = torch.load("last_checkpoint.pt", map_location=DEVICE, weights_only=True)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()
except FileNotFoundError:
    raise FileNotFoundError(f"Checkpoint file not found at last_checkpoint.pt")

def process_audio(audio_input):
    """Process audio to predict phonemes and display mel spectrogram."""
    try:
        print(f"Received audio_input before processing: {audio_input}")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        audio_path = os.path.join(AUDIO_DIR, f"input_{timestamp}.wav")

        if audio_input is None:
            print("Audio input is None")
            return {"error": "No audio input provided"}, None, None, None
        
        if isinstance(audio_input, str):
            print(f"Processing uploaded file: {audio_input}")
            if not os.path.exists(audio_input):
                return {"error": f"Uploaded file not found: {audio_input}"}, None, None, None
            if audio_input.endswith(".mp3"):
                print("Converting .mp3 to .wav")
                from pydub import AudioSegment
                audio = AudioSegment.from_mp3(audio_input)
                audio_path = audio_path.replace(".wav", "_converted.wav")
                audio.export(audio_path, format="wav")
                print(f"Converted file saved to: {audio_path}")
            else:
                shutil.copy(audio_input, audio_path)
                print(f"Copied file to: {audio_path}")
        else:
            raise ValueError("Microphone input not supported in this configuration")

        mel_path = audio_to_mel(audio_path)
        print(f"Generated mel spectrogram: {mel_path}")
        if not os.path.exists(mel_path):
            return {"error": f"Mel spectrogram file not found: {mel_path}"}, None, None, None
        
        mel_image_path = mel_to_image(mel_path)
        print(f"Generated mel spectrogram image: {mel_image_path}")
        if not os.path.exists(mel_image_path):
            return {"error": f"Mel spectrogram image not found: {mel_image_path}"}, None, None, None
        
        mel = np.load(mel_path)
        print(f"Loaded mel spectrogram shape: {mel.shape}")
        mel_tensor = torch.tensor(mel).unsqueeze(0).to(DEVICE)
        mel_lens = torch.tensor([mel.shape[1]]).to(DEVICE)

        with torch.no_grad():
            ph_pred = model(mel_tensor)
            ph_ids = ph_pred.argmax(-1)[0].cpu().numpy()
            print(f"Predicted phoneme IDs: {ph_ids}")

        ph_seq = [id_to_ph[i] for i in ph_ids if i > 0]
        print(f"Raw phonemes: {ph_seq}")

        post_processed = ctc_post_process(ph_seq)
        print(f"Post-processed phonemes: {post_processed}")

        return {
            "audio_path": audio_path,
            "phonemes": " ".join(ph_seq),
            "post_processed_phonemes": " ".join(post_processed)
        }, mel_image_path, " ".join(ph_seq), " ".join(post_processed)
    except Exception as e:
        print(f"Error in process_audio: {str(e)}")
        return {"error": f"Processing failed: {str(e)}"}, None, None, None

# Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Speech to Phonemes Converter")
    gr.Markdown("Upload audio to predict phonemes and display mel spectrogram. Enter text to convert to phonemes.")

    audio_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio (.wav or .mp3)", interactive=True)
    text_input = gr.Textbox(label="Enter Text", placeholder="Type a sentence to convert to phonemes")
    process_button = gr.Button("Process")

    audio_path_output = gr.Textbox(label="Audio Path")
    mel_image = gr.Image(label="Mel Spectrogram", type="filepath")
    raw_phonemes_audio = gr.Textbox(label="Raw Phonemes (Audio)")
    post_processed_phonemes_audio = gr.Textbox(label="Post-Processed Phonemes (Audio)")
    text_phonemes_output = gr.Textbox(label="Phonemes (Text)")

    def process(audio_input, text_input):
        print(f"Processing inputs - Audio: {audio_input}, Text: {text_input}")
        if audio_input:
            audio_result, mel_image_path, raw_ph_audio, post_ph_audio = process_audio(audio_input)
            audio_path = audio_result.get("audio_path", "")
            raw_ph_audio = audio_result.get("phonemes", "")
            post_ph_audio = audio_result.get("post_processed_phonemes", "")
            error_msg = audio_result.get("error", "")
            if error_msg:
                audio_path = error_msg
                raw_ph_audio = ""
                post_ph_audio = ""
                mel_image_path = None
        else:
            audio_path = ""
            raw_ph_audio = ""
            post_ph_audio = ""
            mel_image_path = None
        text_result = text_to_phonemes(text_input) if text_input and text_input.strip() else {}
        text_phonemes = "".join(text_result)
        return audio_path, mel_image_path, raw_ph_audio, post_ph_audio, text_phonemes

    process_button.click(
        fn=process,
        inputs=[audio_input, text_input],
        outputs=[audio_path_output, mel_image, raw_phonemes_audio, post_processed_phonemes_audio, text_phonemes_output]
    )

if __name__ == "__main__":
    iface.launch(debug=True)