import os
import zipfile

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import gradio as gr
import numpy as np
from tensorflow import keras
from PIL import Image

model = keras.models.load_model('mobilenet_ft.keras')
with zipfile.ZipFile('X_test.zip', 'r') as zip_ref:
    with zip_ref.open('X_test.npy') as f:
        X_test = np.load(f)

y_test = np.load('y_test.npy')

def get_random_test_image(correct, total, bot_correct, bot_total):
    index = np.random.randint(0, len(X_test))
    img_arr = X_test[index]

    # Convert to uint8 for display
    if img_arr.max() <= 1.0:
        img_display = (img_arr * 255).astype('uint8')
    else:
        img_display = img_arr.astype('uint8')

    img = Image.fromarray(img_display)

    return img, index, gr.update(interactive=True), gr.update(
        interactive=True), "", "", correct, total, bot_correct, bot_total


def check_guess(guess, current_index, correct, total, bot_correct, bot_total):
    true_label = int(y_test[current_index])
    is_correct = (guess == "Real" and true_label == 1) or (guess == "AI Generated" and true_label == 0)

    # Player score
    total += 1
    if is_correct:
        correct += 1
        result = f"✅ Correct! It was {guess.lower()}!"
    else:
        true_answer = "Real" if true_label == 1 else "AI Generated"
        result = f"❌ Wrong! It was {true_answer}"

    score_text = f"**Your Score:** {correct}/{total} ({100 * correct / total:.1f}%) \n **Bot Score:** {bot_correct}/{bot_total} ({100 * bot_correct / bot_total:.1f}% if bot_total > 0 else 0)" if bot_total > 0 else f"**Your Score:** {correct}/{total} ({100 * correct / total:.1f}%)"

    # Bot prediction
    img = X_test[current_index]
    img_preprocessed = np.expand_dims((img * 255.00), axis=0)
    prediction = model.predict(img_preprocessed, verbose=0)
    binary_prediction = int(prediction > 0.5)

    bot_total += 1
    if binary_prediction == true_label:
        bot_correct += 1

    bot_result = f"🤖 Bot guessed: {'Real' if binary_prediction == 1 else 'AI Generated'} ({f'✓ Correct!' if binary_prediction == true_label else '✗ Wrong'})"

    # Update score display with both scores
    score_text = f"<h3 style='text-align: center;'>Your Score: {correct}/{total} ({100 * correct / total:.1f}%) | Bot Score: {bot_correct}/{bot_total} ({100 * bot_correct / bot_total:.1f}%)</h3>"

    return result, bot_result, gr.update(interactive=False), gr.update(
        interactive=False), score_text, correct, total, bot_correct, bot_total


custom_css = """
.result-box { 
    min-height: 25px; 
}
.textbox_container button,
.textbox button,
.input-container button,
[data-testid="textbox"] button,
textarea ~ button,
input ~ button {
    display: none !important;
    visibility: hidden !important;
    opacity: 0 !important;
    pointer-events: none !important;
}
.podium {
    text-align: center;
    padding: 30px;
    font-size: 20px;
}
.winner {
    font-size: 60px;
    margin: 20px;
}
.scores {
    font-size: 24px;
    margin: 10px;
}
"""


def show_final_results(correct, total, bot_correct, bot_total):
    if total == 0:
        return (gr.update(visible=False), gr.update(visible=True,
                                                   value="<div class='podium'><h2>Play at least one round first! </h2></div>")
                    , 0, 0, 0, 0)

    player_pct = 100 * correct / total
    bot_pct = 100 * bot_correct / total

    if player_pct > bot_pct:
        winner_emoji = "🥇"
        winner_text = "YOU WIN!"
        message = "Congratulations! You beat the AI! 🎉"
    elif bot_pct > player_pct:
        winner_emoji = "🤖"
        winner_text = "BOT WINS!"
        message = "The robots are getting smarter... 😨 "
    else:
        winner_emoji = "🤝"
        winner_text = "TIE GAME!"
        message = "Perfectly balanced...as all things should be ⚖️"

    results_html = f"""
    <div class='podium'>
        <div class='winner'>{winner_emoji}</div>
        <h1>{winner_text}</h1>
        <p style='font-size: 18px; color: #666;'>{message}</p>
        <hr style='margin: 30px 0;'>
        <div class='scores'>
            <p>👤 <strong>Your Score:</strong> {correct}/{total} ({player_pct:.1f}%)</p>
            <p>🤖 <strong>Bot Score:</strong> {bot_correct}/{total} ({bot_pct:.1f}%)</p>
        </div>
    </div>
    """

    return gr.update(visible=False), gr.update(visible=True, value=results_html), 0, 0, 0, 0

def restart_game():
    return gr.update(visible=True), gr.update(visible=False), 0, 0, 0, 0

with (gr.Blocks(theme=gr.themes.Glass(), css=custom_css)
      as demo):
    gr.Markdown("<h1 style='text-align: center;'>Man vs. Bot</h1>")
    gr.Markdown("<h3 style='text-align: center;'>Is this image real or AI-generated?</h3>")

    game_interface = gr.Column(visible=True)
    results_interface = gr.Markdown(visible=False)

    current_index = gr.State(0)
    correct_count = gr.State(0)
    total_count = gr.State(0)
    bot_correct_count = gr.State(0)
    bot_total_count = gr.State(0)

    with game_interface:
        output_image = gr.Image(
            label="Random Test Image",
            show_label=False,
            container=False,
            show_download_button=False,
            show_share_button=False
        )

        with gr.Row():
            real_btn = gr.Button("Real", variant="primary")
            ai_btn = gr.Button("AI Generated", variant="primary")

        with gr.Row():
            result_text = gr.Textbox(
                label="",
                interactive=False,
                show_label=False,
                show_copy_button=False,
                container=False
            )
            bot_result_text = gr.Textbox(
                label="",
                interactive=False,
                show_label=False,
                show_copy_button=False,
                container=False
            )

        score_display = gr.Markdown("<h3 style='text-align: center;'> 🧑‍🦲 Your Score: 0/0 | 🤖 Bot Score: 0/0</h3>")
        next_btn = gr.Button("Next", variant="primary")
        quit_btn = gr.Button("Quit")

    restart_btn = gr.Button("Play Again", visible=False)

    demo.load(fn=get_random_test_image, inputs=[correct_count, total_count, bot_correct_count, bot_total_count],
              outputs=[output_image, current_index, real_btn, ai_btn, result_text,
                       bot_result_text, correct_count,
                       total_count, bot_correct_count, bot_total_count])

    real_btn.click(
        fn=lambda idx, c, t, bc, bt: check_guess("Real", idx, c, t, bc, bt),
        inputs=[current_index, correct_count, total_count, bot_correct_count, bot_total_count],
        outputs=[result_text, bot_result_text, real_btn, ai_btn, score_display,
                 correct_count, total_count,
                 bot_correct_count, bot_total_count]
    )
    ai_btn.click(
        fn=lambda idx, c, t, bc, bt: check_guess("AI Generated", idx, c, t, bc, bt),
        inputs=[current_index, correct_count, total_count, bot_correct_count, bot_total_count],
        outputs=[result_text, bot_result_text, real_btn, ai_btn, score_display,
                 correct_count, total_count,
                 bot_correct_count, bot_total_count]
    )

    next_btn.click(
        fn=get_random_test_image,
        inputs=[correct_count, total_count, bot_correct_count, bot_total_count],
        outputs=[output_image, current_index, real_btn, ai_btn, result_text, bot_result_text, correct_count,
                 total_count, bot_correct_count, bot_total_count]
    )

    quit_btn.click(
        fn=show_final_results,
        inputs=[correct_count, total_count, bot_correct_count, bot_total_count],
        outputs=[game_interface, results_interface, correct_count, total_count,
                 bot_correct_count, bot_total_count]
    ).then(
        fn=lambda: gr.update(visible=True),
        outputs=[restart_btn]
    )

    restart_btn.click(
        fn=restart_game,
        outputs=[game_interface, results_interface, correct_count, total_count,
                 bot_correct_count, bot_total_count]
    ).then(
        fn=lambda: "<h3 style='text-align: center;'>Your Score: 0/0 | Bot Score: 0/0</h3>",
        outputs=[score_display]
    ).then(
        fn=lambda: gr.update(visible=False),
        outputs=[restart_btn]
    ).then(
        fn=get_random_test_image,
        inputs=[correct_count, total_count, bot_correct_count, bot_total_count],
        outputs=[output_image, current_index, real_btn, ai_btn, result_text,
                 bot_result_text, correct_count,
                 total_count, bot_correct_count, bot_total_count]
    )

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