import logging

import gradio as gr

from constants import ALL_PHONEMES, MODEL_REPO_ID
from utils import load_model_and_processor, parse_delta_value, run_inference, validate_phonemes

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)

logger.info("Loading model into memory globally")
model, processor = load_model_and_processor(MODEL_REPO_ID)
logger.info("Model loaded successfully")

css = """
.phoneme-scores { display: flex; flex-wrap: wrap; justify-content: center; gap: 15px; }
.phoneme-container { text-align: center; padding: 10px; border: 1px solid #ddd; border-radius: 8px; }
.phoneme { font-size: 1.5em; font-weight: bold; margin-bottom: 5px; }
.score { padding: 8px 12px; border-radius: 5px; color: white; font-weight: bold; }
.good { background-color: #28a745; }
.medium { background-color: #ffc107; }
.bad { background-color: #dc3545; }
"""


def get_score_class(score, score_type):
    if score_type == "quality":
        if score == 1:
            return "good"
        if score == 2:
            return "medium"
        return "bad"
    return "good" if score == 1 else "bad"


def generate_html_output(result, score_type):
    if isinstance(result, str):
        return result

    if not result:
        return "<p style='text-align:center; color:red;'>No scores were produced.</p>"

    predicted_scores, tokens = result
    title = "Quality Scores" if score_type == "quality" else "Duration Scores"
    html_output = f"<div class='phoneme-section'><h3 class='scores-title'>{title}</h3></div><div class='phoneme-scores'>"

    for token, score in zip(tokens, predicted_scores):
        display_score = int(score) + 1
        score_class = get_score_class(display_score, score_type)
        html_output += f"""
        <div class='phoneme-container'>
            <div class='phoneme'>{token}</div>
            <div class='score {score_class}'>{display_score}</div>
        </div>
        """

    html_output += "</div>"
    return html_output


DEFAULT_QUALITY_DELTA = 0.56
DEFAULT_DURATION_DELTA = 0.5


def score_phonemes(phoneme_text, audio_file, quality_delta, duration_delta):
    if audio_file is None:
        return "<p style='text-align:center; color:red;'>Please upload a .wav audio file.</p>", ""

    phonemes_validation_error = validate_phonemes(phoneme_text, ALL_PHONEMES)
    if phonemes_validation_error:
        return phonemes_validation_error, ""

    results = run_inference(
        audio_file,
        phoneme_text,
        model,
        processor,
        deltas={
            "quality": parse_delta_value(quality_delta),
            "duration": parse_delta_value(duration_delta),
        },
        correct_index=0,
    )
    if isinstance(results, str):
        return results, ""

    quality_result = results.get("quality")
    duration_result = results.get("duration")

    quality_html = generate_html_output(quality_result, "quality")
    duration_html = generate_html_output(duration_result, "duration")

    return quality_html, duration_html


def score_phonemes_json(phoneme_text, audio_file):
    if audio_file is None:
        return None

    phonemes_validation_error = validate_phonemes(phoneme_text, ALL_PHONEMES)
    if phonemes_validation_error:
        return None

    results = run_inference(
        audio_file,
        phoneme_text,
        model,
        processor,
        deltas={
            "quality": DEFAULT_QUALITY_DELTA,
            "duration": DEFAULT_DURATION_DELTA,
        },
        correct_index=0,
    )
    if isinstance(results, str) or not results:
        return None

    quality_result = results.get("quality")
    if isinstance(quality_result, str) or not quality_result:
        return None

    scores, tokens = quality_result
    return [{"token": t, "quality": int(s) + 1} for t, s in zip(tokens, scores)]


with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Phoneme Pronunciation and Duration Scorer
        Enter phonemes directly into the text box, separated by spaces.
        Use `|` between words if you want to score a multi-word sequence.
        Then upload a `.wav` file or record the audio of the pronounced word.
        The application will provide a pronunciation (quality) and duration score for each phoneme.
        Scores legend:
        - Quality: 1 (good), 2 (medium), 3 (bad)
        - Duration: 1 (good), 2 (bad)
        """
    )

    with gr.Row():
        word_input = gr.Textbox(label="Word")
        phoneme_text_input = gr.Textbox(label="Phonemes (space-separated)")
        audio_input = gr.Audio(type="filepath", label="Audio File (.wav)")

    with gr.Row():
        quality_delta_input = gr.Number(label="Quality Delta", value=DEFAULT_QUALITY_DELTA, precision=2)
        duration_delta_input = gr.Number(label="Duration Delta", value=DEFAULT_DURATION_DELTA, precision=2)

    btn = gr.Button("Generate Scores", variant="primary")

    gr.Examples(
        examples=[
            ["aia (L1 speaker)", "a i j a", "./audio/L1/e7cd-68c0-b5df-35b0_aia_take1.wav"],
            ["maias (L1 speaker)", "m a i j a s", "./audio/L1/e7cd-68c0-b5df-35b0_maias_take1.wav"],
            ["aja (L2 speaker)", "a j a", "./audio/L2/03ac-e45b-ec8a-6fa0_aja_take1.wav"],
            ["sõpra (L2 speaker)", "s õ pp r a", "./audio/L2/4071-0c77-e1d3-9587_sõpra_take1.wav"],
        ],
        inputs=[word_input, phoneme_text_input, audio_input],
    )

    gr.Markdown("---")
    gr.Markdown("## Results")

    with gr.Row():
        phoneme_output_html = gr.HTML()
        duration_output_html = gr.HTML()

    btn.click(
        fn=score_phonemes,
        inputs=[phoneme_text_input, audio_input, quality_delta_input, duration_delta_input],
        outputs=[phoneme_output_html, duration_output_html],
        api_name="L2_verifier",
    )

    # Hidden components for the JSON API endpoint
    json_phoneme = gr.Textbox(visible=False)
    json_audio = gr.Audio(type="filepath", visible=False)
    json_output = gr.JSON(visible=False)
    json_btn = gr.Button(visible=False)

    json_btn.click(
        fn=score_phonemes_json,
        inputs=[json_phoneme, json_audio],
        outputs=[json_output],
        api_name="produce_scores",
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch()