import os

# ⚙️ Ensure model is saved before import
if not os.path.exists("trained_model/config.json"):
    print("⚙️ No model found. Initializing and saving EvoTransformer...")
    from init_model import initialize_and_save_model
    initialize_and_save_model()

import gradio as gr
import random
from inference import generate_response
from logger import log_user_feedback
from dashboard import (
    update_dashboard_plot,
    evolution_accuracy_plot,
    leaderboard_plot,
    get_vote_counts,
)
from watchdog import retrain_model
from init_model import load_model

model = load_model()

def get_architecture_summary(model):
    return "\n".join([
        f"Layers: {getattr(model, 'num_layers', 'N/A')}",
        f"Attention Heads: {getattr(model, 'num_heads', 'N/A')}",
        f"FFN Dim: {getattr(model, 'ffn_dim', 'N/A')}",
        f"Memory Enabled: {getattr(model, 'use_memory', 'N/A')}",
    ])

def update_vote_counter():
    votes = get_vote_counts()
    s1 = votes.get("1", 0)
    s2 = votes.get("2", 0)
    total = s1 + s2
    return f"🗳️ **Votes so far** — Solution 1: {s1} | Solution 2: {s2} | Total: {total}"

examples = [
    {"goal": "Escape from a burning house", "option1": "Run out through the front door", "option2": "Hide in the bathroom"},
    {"goal": "Improve sleep quality", "option1": "Use phone in bed", "option2": "Turn off screens 1 hour before bed"},
    {"goal": "Lose weight", "option1": "Skip meals", "option2": "Exercise and eat healthy"},
]

def load_random_example():
    example = random.choice(examples)
    return example["goal"], example["option1"], example["option2"]

def evo_chat(goal, sol1, sol2):
    result = generate_response(goal, sol1, sol2)
    return result["evo_suggestion"], result["gpt_suggestion"]

def handle_feedback(goal, sol1, sol2, winner):
    try:
        log_user_feedback(goal, sol1, sol2, winner)
        return "✅ Feedback logged. Thank you!"
    except Exception as e:
        return f"❌ Failed to log: {e}"

with gr.Blocks(title="EvoTransformer v2.1") as demo:
    gr.Markdown("""
# EvoTransformer v2.1  
*Built Different. Learns Live.*

🧬 Learns from your input  
🧠 Compares reasoning like GPT-3.5  
📊 Tracks accuracy & evolves
""")

    with gr.Row():
        goal = gr.Textbox(label="Goal", placeholder="e.g. Escape from house on fire")
    with gr.Row():
        sol1 = gr.Textbox(label="Option 1", placeholder="e.g. Exit through door")
        sol2 = gr.Textbox(label="Option 2", placeholder="e.g. Hide under bed")

    with gr.Row():
        compare_btn = gr.Button("🔍 Compare")
        random_btn = gr.Button("🎲 Load Random Example")

    evo_box = gr.Textbox(label="🧠 Evo Suggestion")
    gpt_box = gr.Textbox(label="💬 GPT-3.5 Suggestion")
    feedback_msg = gr.Textbox(visible=False)

    with gr.Row():
        vote = gr.Radio(["Solution 1", "Solution 2"], label="Which was better?")
        log_btn = gr.Button("✅ Log Feedback")

    compare_btn.click(fn=evo_chat, inputs=[goal, sol1, sol2], outputs=[evo_box, gpt_box])
    random_btn.click(fn=load_random_example, outputs=[goal, sol1, sol2])
    log_btn.click(fn=handle_feedback, inputs=[goal, sol1, sol2, vote], outputs=[feedback_msg])

    # === Visual Stats ===
    vote_counter = gr.Markdown(update_vote_counter())
    arch_summary = gr.Textbox(value=get_architecture_summary(model), label="📐 Model Configuration", lines=4, interactive=False)
    evo_plot = gr.Image(value=evolution_accuracy_plot(), label="📈 Accuracy Over Time")
    dash_img = gr.Image(value=update_dashboard_plot(), label="📊 User Feedback Distribution")
    leaderboard = gr.Image(value=leaderboard_plot(), label="🏆 Hall of Fame")

    retrain_note = gr.Markdown("Click to retrain Evo using your latest feedback.")
    retrain_status = gr.Textbox(label="Retrain Status", value="No retrain triggered yet.")
    retrain_btn = gr.Button("♻️ Retrain Evo", variant="primary")
    retrain_btn.click(fn=retrain_model, outputs=[arch_summary, evo_plot, retrain_status])

    with gr.Accordion("🧬 EvoTransformer Architecture", open=False): arch_summary
    with gr.Accordion("📈 Accuracy Plot", open=False): evo_plot
    with gr.Accordion("📊 Vote Summary", open=False): dash_img
    with gr.Accordion("🏅 Top Goals", open=False): leaderboard
    with gr.Accordion("🔁 Retrain Log", open=False): retrain_status

    gr.Markdown("Built with ❤️ using Gradio · Contact us for API access")

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