import gradio as gr
import os
import json
from datetime import datetime
from pathlib import Path
import soundfile as sf
from core.cloner import KokoClone

import whisper
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# === MODELOS BASE ===
print("Loading Whisper + Translator models...")
whisper_model = whisper.load_model("small")
translator_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-es-en")
translator_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-es-en")

print("Loading KokoClone models for the Web UI...")
cloner = KokoClone()

# === MODELO DE CORRECCIÓN GRAMATICAL BRITÁNICA ===
grammar_tokenizer = AutoTokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
grammar_model = AutoModelForSeq2SeqLM.from_pretrained("prithivida/grammar_error_correcter_v1")

# === Memoria de conversación (solo durante la sesión) ===
conversation_history = []

# === Diccionario personal ===
DATA_DIR = Path("/app/user_data")
AUDIO_DIR = DATA_DIR / "audios"
DICT_PATH = DATA_DIR / "dictionary.json"

DATA_DIR.mkdir(parents=True, exist_ok=True)
AUDIO_DIR.mkdir(parents=True, exist_ok=True)

def load_dictionary():
    if DICT_PATH.exists():
        with open(DICT_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    return []

def save_dictionary(entries):
    with open(DICT_PATH, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

def add_entry(entry):
    entries = load_dictionary()
    entries.append(entry)
    save_dictionary(entries)

# === Spanish speech → English text ===
def spanish_to_english(audio_file):
    result = whisper_model.transcribe(audio_file)
    spanish_text = result["text"]

    inputs = translator_tokenizer(spanish_text, return_tensors="pt")
    translated_tokens = translator_model.generate(**inputs)
    english_text = translator_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)

    return spanish_text, english_text

# === Clone text → voice ===
def clone_voice(text, lang, ref_audio_path):
    if not text or not text.strip():
        raise gr.Error("Please enter some text.")
    if not ref_audio_path:
        raise gr.Error("Please upload or record a reference audio file.")

    output_file = "gradio_output.wav"

    try:
        cloner.generate(
            text=text,
            lang=lang,
            reference_audio=ref_audio_path,
            output_path=output_file
        )
        return output_file
    except Exception as e:
        raise gr.Error(f"An error occurred during generation: {str(e)}")

# === Convert audio → revoice ===
def convert_voice(source_audio_path, ref_audio_path):
    if not source_audio_path:
        raise gr.Error("Please upload or record a source audio file.")
    if not ref_audio_path:
        raise gr.Error("Please upload or record a reference audio file.")

    output_file = "gradio_convert_output.wav"

    try:
        cloner.convert(
            source_audio=source_audio_path,
            reference_audio=ref_audio_path,
            output_path=output_file
        )
        return output_file
    except Exception as e:
        raise gr.Error(f"An error occurred during conversion: {str(e)}")

# === Spanish speech → English clone + dictionary ===
def mic_to_english_clone(audio_file, ref_audio_path):
    if not audio_file:
        raise gr.Error("Please speak into the microphone.")
    if not ref_audio_path:
        raise gr.Error("Please upload a reference audio (your voice).")

    spanish_text, english_text = spanish_to_english(audio_file)

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    output_file = AUDIO_DIR / f"mic_english_clone_{timestamp}.wav"

    cloner.generate(
        text=english_text,
        lang="en",
        reference_audio=ref_audio_path,
        output_path=str(output_file)
    )

    entry = {
        "fecha": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "modo": "hablar_espanol_a_ingles",
        "espanol": spanish_text,
        "ingles_traducido": english_text,
        "audio_clonado": str(output_file),
        "nota": "shadowing"
    }
    add_entry(entry)

    return str(output_file), english_text

# === Pronunciación → corrección (profesor británico estricto) ===
def pronunciation_correction(audio_file, target_text, ref_audio_path):
    if not audio_file:
        raise gr.Error("Please speak the sentence in English.")
    if not target_text or not target_text.strip():
        raise gr.Error("Please provide the target English sentence.")
    if not ref_audio_path:
        raise gr.Error("Please upload a reference audio (your voice).")

    result = whisper_model.transcribe(audio_file)
    spoken_text = result["text"].strip()

    target_clean = target_text.strip()
    spoken_clean = spoken_text.strip()

    target_words = target_clean.lower().split()
    spoken_words = spoken_clean.lower().split()

    matches = sum(1 for w in spoken_words if w in target_words)
    score = int(100 * matches / max(len(target_words), 1))

    feedback_lines = []
    feedback_lines.append("British strict professor report")
    feedback_lines.append(f"Target sentence: {target_clean}")
    feedback_lines.append(f"You said:       {spoken_clean}")
    feedback_lines.append(f"Pronunciation score: {score}%")

    if score == 100:
        feedback_lines.append("Excellent. Your pronunciation and wording are accurate. We may proceed.")
    elif score >= 80:
        feedback_lines.append("Good attempt, but not perfect. Some words were unclear or slightly off.")
        feedback_lines.append("Repeat the sentence focusing on rhythm and clarity.")
    else:
        feedback_lines.append("Your sentence is not sufficiently accurate.")
        feedback_lines.append("You must repeat it precisely before we continue.")
        feedback_lines.append("Listen carefully to the corrected version and shadow it several times.")

    feedback = "\n".join(feedback_lines)

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    output_file = AUDIO_DIR / f"pronunciation_correct_{timestamp}.wav"

    cloner.generate(
        text=target_clean,
        lang="en",
        reference_audio=ref_audio_path,
        output_path=str(output_file)
    )

    entry = {
        "fecha": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "modo": "pronunciacion_profesor_britanico_estricto",
        "frase_objetivo": target_clean,
        "frase_dicha": spoken_clean,
        "score": score,
        "audio_clonado": str(output_file),
        "feedback": feedback
    }
    add_entry(entry)

    return str(output_file), feedback

# === CORRECCIÓN GRAMATICAL BRITÁNICA ===
def british_grammar_fix(sentence):
    inputs = grammar_tokenizer.encode(sentence, return_tensors="pt")
    outputs = grammar_model.generate(inputs, max_length=128)
    corrected = grammar_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected

# === CONVERSACIÓN AVANZADA (PROFESOR LONDINENSE NATURAL) ===
def conversation_professor(audio_file, ref_audio_path, topic):
    if not audio_file:
        raise gr.Error("Please speak into the microphone.")
    if not ref_audio_path:
        raise gr.Error("Please upload your reference voice.")
    if not topic:
        raise gr.Error("Please choose a topic for the conversation.")

    # 1. Transcribir lo que has dicho
    result = whisper_model.transcribe(audio_file)
    spoken_text = result["text"].strip()

    # 2. Corregir gramática británica
    corrected = british_grammar_fix(spoken_text)

    # 3. Guardar en historial
    global conversation_history
    conversation_history.append({
        "user_raw": spoken_text,
        "user_corrected": corrected,
        "topic": topic,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    })

    # 4. Generar feedback estilo londinense natural
    feedback_lines = []
    feedback_lines.append(f"Topic: {topic}")
    feedback_lines.append(f"You said: {spoken_text}")
    feedback_lines.append(f"Natural British version: {corrected}")
    feedback_lines.append("")

    if corrected.lower() == spoken_text.lower():
        feedback_lines.append(
            "Nice one. That already sounds quite natural. "
            "Next time, try to expand your answer a bit more."
        )
    else:
        feedback_lines.append(
            "Not bad, but there are a few things to polish. "
            "Use the corrected version and say it again, a bit more confidently."
        )

    if len(conversation_history) > 1:
        last = conversation_history[-2]
        feedback_lines.append("")
        feedback_lines.append(
            "Earlier you said: \"" + last["user_corrected"] + "\" "
            "— you're slowly building a nice flow. Keep going."
        )

    feedback_lines.append("")
    if topic == "Restaurant":
        feedback_lines.append(
            "Imagine we're still at the restaurant in London. "
            "Now tell me what you would like to order, using at least two adjectives."
        )
    elif topic == "Hotel":
        feedback_lines.append(
            "We're at the hotel reception. "
            "Now explain politely that you need a late check-out tomorrow."
        )
    elif topic == "Airport":
        feedback_lines.append(
            "We're at Heathrow. "
            "Now tell me what you would say if your flight was delayed by three hours."
        )
    elif topic == "Work":
        feedback_lines.append(
            "We're at the office. "
            "Now describe your typical working day in three sentences."
        )
    else:
        feedback_lines.append(
            "Carry on talking about this topic, but try to be a bit more specific."
        )

    feedback = "\n".join(feedback_lines)

    # 5. Generar audio con la versión corregida (tu voz clonada)
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    output_file = AUDIO_DIR / f"professor_conversation_{timestamp}.wav"

    cloner.generate(
        text=corrected,
        lang="en",
        reference_audio=ref_audio_path,
        output_path=str(output_file)
    )

    return str(output_file), corrected, feedback

# === UI ===
with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div style="text-align: center;">
            <h1>Miguel KokoClone</h1>
            <p>Voice Cloning + English Trainer (British strict professor + London conversation).<br>
            Daily shadowing, correction and personal dictionary.</p>
        </div>
        """
    )

    with gr.Tabs():

        # Text → Clone
        with gr.Tab("Text → Clone"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="1. Text to Synthesize",
                        lines=4,
                        placeholder="Enter the text you want spoken..."
                    )

                    lang_input = gr.Dropdown(
                        label="2. Language",
                        choices=[
                            ("English", "en"),
                            ("Hindi", "hi"),
                            ("French", "fr"),
                            ("Japanese", "ja"),
                            ("Chinese", "zh"),
                            ("Italian", "it"),
                            ("Spanish", "es"),
                            ("Portuguese", "pt")
                        ],
                        value="en"
                    )

                    ref_audio_input = gr.Audio(
                        label="3. Reference Voice (Upload or Record)",
                        type="filepath"
                    )

                    submit_btn = gr.Button("Generate Clone", variant="primary")

                with gr.Column():
                    output_audio = gr.Audio(
                        label="Generated Cloned Audio",
                        interactive=False,
                        autoplay=False
                    )

            submit_btn.click(
                fn=clone_voice,
                inputs=[text_input, lang_input, ref_audio_input],
                outputs=output_audio
            )

        # Audio → Clone
        with gr.Tab("Audio → Clone"):
            with gr.Row():
                with gr.Column():
                    source_audio_input = gr.Audio(
                        label="1. Source Audio (speech to re-voice)",
                        type="filepath"
                    )

                    ref_audio_convert_input = gr.Audio(
                        label="2. Reference Voice (target speaker)",
                        type="filepath"
                    )

                    convert_btn = gr.Button("Convert Voice", variant="primary")

                with gr.Column():
                    convert_output_audio = gr.Audio(
                        label="Converted Audio",
                        interactive=False,
                        autoplay=False
                    )

            convert_btn.click(
                fn=convert_voice,
                inputs=[source_audio_input, ref_audio_convert_input],
                outputs=convert_output_audio
            )

        # Hablar → Inglés (shadowing + diccionario)
        with gr.Tab("Hablar → Inglés (voz clonada + texto)"):
            with gr.Row():
                with gr.Column():
                    mic_input = gr.Audio(
                        label="1. Speak in Spanish (microphone)",
                        sources=["microphone"],
                        type="filepath"
                    )

                    ref_audio_mic_input = gr.Audio(
                        label="2. Reference Voice (your voice)",
                        type="filepath"
                    )

                    mic_btn = gr.Button("Speak → English Clone", variant="primary")

                with gr.Column():
                    mic_output_audio = gr.Audio(
                        label="English Cloned Audio",
                        interactive=False,
                        autoplay=False
                    )

                    english_text_output = gr.Textbox(
                        label="English Transcription",
                        lines=3,
                        interactive=False
                    )

            mic_btn.click(
                fn=mic_to_english_clone,
                inputs=[mic_input, ref_audio_mic_input],
                outputs=[mic_output_audio, english_text_output]
            )

        # Pronunciación → corrección (profesor británico estricto)
        with gr.Tab("Pronunciación (profesor británico estricto)"):
            with gr.Row():
                with gr.Column():
                    pron_audio = gr.Audio(
                        label="1. Say the sentence in English (microphone)",
                        sources=["microphone"],
                        type="filepath"
                    )

                    target_sentence = gr.Textbox(
                        label="2. Target English sentence (British)",
                        lines=2,
                        placeholder="Type the exact sentence you want to pronounce..."
                    )

                    ref_audio_pron = gr.Audio(
                        label="3. Reference Voice (your voice)",
                        type="filepath"
                    )

                    pron_btn = gr.Button("Evaluate pronunciation", variant="primary")

                with gr.Column():
                    pron_output_audio = gr.Audio(
                        label="Correct British version (your cloned voice)",
                        interactive=False,
                        autoplay=False
                    )

                    pron_feedback = gr.Textbox(
                        label="Professor feedback",
                        lines=8,
                        interactive=False
                    )

            pron_btn.click(
                fn=pronunciation_correction,
                inputs=[pron_audio, target_sentence, ref_audio_pron],
                outputs=[pron_output_audio, pron_feedback]
            )

        # Conversación avanzada (profesor londinense)
        with gr.Tab("Conversación avanzada (profesor londinense)"):
            with gr.Row():
                with gr.Column():
                    conv_topic = gr.Dropdown(
                        label="Conversation topic",
                        choices=[
                            "Restaurant",
                            "Hotel",
                            "Airport",
                            "Work",
                            "Casual"
                        ],
                        value="Casual"
                    )

                    conv_audio = gr.Audio(
                        label="Speak in English (microphone)",
                        sources=["microphone"],
                        type="filepath"
                    )
                    conv_ref = gr.Audio(
                        label="Reference Voice (your voice)",
                        type="filepath"
                    )
                    conv_btn = gr.Button("Send to professor", variant="primary")

                with gr.Column():
                    conv_output_audio = gr.Audio(
                        label="British corrected version (your cloned voice)",
                        interactive=False,
                        autoplay=False
                    )
                    conv_corrected_text = gr.Textbox(
                        label="Natural British sentence",
                        lines=3,
                        interactive=False
                    )
                    conv_feedback = gr.Textbox(
                        label="Professor (London) feedback + next question",
                        lines=10,
                        interactive=False
                    )

            conv_btn.click(
                fn=conversation_professor,
                inputs=[conv_audio, conv_ref, conv_topic],
                outputs=[conv_output_audio, conv_corrected_text, conv_feedback]
            )

        # Mi diccionario
        with gr.Tab("Mi diccionario"):
            dict_view = gr.JSON(
                label="Stored entries (Spanish → English, pronunciation, professor feedback)",
                value=load_dictionary()
            )

# === Launch ===
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", share=True)