import gradio as gr
import os
from huggingface_hub import InferenceClient, __version__ as hf_version
import random
from typing import Generator, Dict, List, Tuple, Optional
import logging

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logging.debug(f"Using huggingface_hub version: {hf_version}")

# Get token from environment variable
hf_token = os.environ.get("QWEN_BOT_TOKEN")
client = InferenceClient("Qwen/Qwen2.5-72B-Instruct", token=hf_token)

TOPIC_EXAMPLES = {
    "Daily Life": {
        "beginner": [
            "What time do you wake up?",
            "Do you go to school or work?",
            "What do you eat for breakfast?"
        ],
        "intermediate": [
            "What do you usually do after work?",
            "Do you like cooking? What’s your favorite dish?",
            "Tell me about your morning routine."
        ],
        "advanced": [
            "How do you balance personal and professional responsibilities?",
            "What does a productive day look like for you?",
            "How has your daily routine changed over the years?"
        ]
    },
    "Travel": {
        "beginner": [
            "Do you like to travel?",
            "Have you been to another city?",
            "Do you like airplanes?"
        ],
        "intermediate": [
            "Have you ever been to another country?",
            "What's your dream vacation?",
            "What do you like to pack in your suitcase?"
        ],
        "advanced": [
            "How has travel influenced your worldview?",
            "What do you think makes a destination culturally significant?",
            "Describe your most challenging travel experience."
        ]
    },
    "Food": {
        "beginner": [
            "Do you like apples?",
            "What is your favorite snack?",
            "Do you eat rice or noodles?"
        ],
        "intermediate": [
            "Can you describe how to cook your favorite dish?",
            "What is your go-to comfort food and why?",
            "Have you ever tried food from another country?"
        ],
        "advanced": [
            "How does food reflect a culture’s values and traditions?",
            "Compare the cuisines of two different countries.",
            "What’s the most unique food you’ve ever tasted?"
        ]
    },
    "Work & School": {
        "beginner": [
            "Do you go to school or work?",
            "What is your teacher’s name?",
            "Do you have homework?"
        ],
        "intermediate": [
            "What do you do at your job or school?",
            "What’s your favorite subject or task?",
            "Do you like working with other people?"
        ],
        "advanced": [
            "How do you stay motivated in your work or studies?",
            "What are the challenges of remote learning or working?",
            "How do education systems differ around the world?"
        ]
    },
    "Hobbies": {
        "beginner": [
            "Do you like music?",
            "What games do you play?",
            "Can you draw or paint?"
        ],
        "intermediate": [
            "What hobbies do you enjoy in your free time?",
            "When did you start your favorite hobby?",
            "Do you prefer indoor or outdoor hobbies?"
        ],
        "advanced": [
            "How can hobbies contribute to personal growth?",
            "What’s a hobby you would like to master and why?",
            "How has technology changed the way we pursue hobbies?"
        ]
    },
    "Shopping": {
        "beginner": [
            "Do you like shopping?",
            "What do you buy at the store?",
            "Do you go shopping alone?"
        ],
        "intermediate": [
            "Do you prefer shopping online or in stores?",
            "Tell me about your last shopping trip.",
            "What kinds of things do you usually buy?"
        ],
        "advanced": [
            "How has consumer behavior changed over time?",
            "What are the pros and cons of online shopping?",
            "Do advertisements affect your shopping decisions?"
        ]
    },
    "Weather": {
        "beginner": [
            "Is it sunny today?",
            "Do you like rain?",
            "What is your favorite season?"
        ],
        "intermediate": [
            "What’s the weather like where you are?",
            "Do you like hot or cold weather?",
            "How do you prepare for a rainy day?"
        ],
        "advanced": [
            "How does climate affect daily life in your region?",
            "What are the consequences of global climate change?",
            "How does weather influence culture and traditions?"
        ]
    }
}

MAX_HISTORY_LENGTH = 20
MEMORY_WINDOW = 5
MAX_TOKENS = 1024
TEMPERATURE = 0.7
TOP_P = 0.95

def get_examples_for_topic(topic, difficulty):
    return TOPIC_EXAMPLES.get(topic, {}).get(difficulty, [])

def get_conversation_prompt():
    return """You are JoJo, a friendly and supportive AI who helps people practice speaking English through fun, casual conversation.

Always keep your replies short (2–4 sentences) and speak naturally, like a friendly tutor or partner. Ask engaging, open-ended questions to keep the conversation going.

If the user makes a grammar or vocabulary mistake, gently correct them by repeating the corrected sentence, followed by a kind note like: \"You can also say it like this!\"

Do not teach grammar rules. Just offer corrections through example. Never lecture or explain unless asked.

Keep the tone warm, encouraging, and non-judgmental."""

def respond(message: str, chat_history: List[Tuple[str, str]], topic: Optional[str] = None, use_full_memory: bool = True) -> Tuple[str, List[Tuple[str, str]]]:
    if not message.strip():
        return "", chat_history

    try:
        api_messages = [{"role": "system", "content": get_conversation_prompt()}]
        logging.debug(f"System Message: {api_messages[0]}")

        if chat_history and use_full_memory:
            for user_msg, bot_msg in chat_history[-MEMORY_WINDOW:]:
                api_messages.extend([
                    {"role": "user", "content": str(user_msg)},
                    {"role": "assistant", "content": str(bot_msg)}
                ])

        api_messages.append({"role": "user", "content": str(message)})
        response = client.chat_completion(
            messages=api_messages,
            max_tokens=MAX_TOKENS,
            temperature=TEMPERATURE,
            top_p=TOP_P
        )

        bot_message = response.choices[0].message.content
        updated_history = list(chat_history)
        updated_history.append((message, bot_message))
        return "", updated_history

    except Exception as e:
        logging.error("Error in respond function", exc_info=True)
        error_msg = f"There was a temporary issue. Please try again. (Error: {str(e)})"
        return "", list(chat_history) + [(message, error_msg)]

def get_avatar_url():
    return "https://api.dicebear.com/7.x/bottts/svg?seed=rabbit&backgroundColor=b6e3f4"

custom_css = """
.compact-btn {
    padding: 0.75rem !important;
    font-size: 1rem !important;
    font-weight: 500;
    border-radius: 8px;
    background-color: #2f2f2f;
    color: white;
    transition: background-color 0.3s;
}
.compact-btn:hover {
    background-color: #444;
}
#voice-controls {
    margin-top: 1em;
    text-align: center;
    opacity: 0.5;
    font-size: 0.9rem;
    font-style: italic;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("""
    # 🐰 JoJo - Your Speaking Buddy
    **Chat in English with JoJo — your kind and cheerful language partner.**  
    Pick a topic, choose your level, and practice naturally. JoJo will guide you, ask questions, and gently correct you along the way!
    """)

    avatar = get_avatar_url()
    memory_flag = gr.State(value=True)

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                height=400,
                bubble_full_width=True,
                show_copy_button=True,
                avatar_images=[None, avatar],
                scale=1,
                min_width=800
            )
            msg = gr.Textbox(
                placeholder="Say something to JoJo...",
                scale=4
            )

            with gr.Row():
                submit = gr.Button("Send", variant="primary")
                clear = gr.Button("New Chat")

            gr.Markdown("""
            <div id="voice-controls">
            🎤 Voice input and 🔈 playback coming soon!
            </div>
            """)

        with gr.Column(scale=1):
            gr.Markdown("""### 🎯 Conversation Settings""")
            topic = gr.Dropdown(
                choices=list(TOPIC_EXAMPLES.keys()),
                label="Select Topic",
                value="Daily Life"
            )

            difficulty = gr.Dropdown(
                choices=["beginner", "intermediate", "advanced"],
                label="Select Difficulty",
                value="intermediate"
            )

            gr.Markdown("""### 💬 Quick Starters""")
            starter_btn1 = gr.Button("Starter 1", scale=1, min_width=250, elem_classes="compact-btn")
            starter_btn2 = gr.Button("Starter 2", scale=1, min_width=250, elem_classes="compact-btn")
            starter_btn3 = gr.Button("Starter 3", scale=1, min_width=250, elem_classes="compact-btn")
            starter_buttons = [starter_btn1, starter_btn2, starter_btn3]

            def update_starters(selected_topic, selected_difficulty):
                examples = get_examples_for_topic(selected_topic, selected_difficulty)
                results = [examples[i] if i < len(examples) else "" for i in range(3)]
                return tuple(results)

            def use_starter(text: str, history: List[Tuple[str, str]], selected_topic: str, memory_flag: bool) -> Tuple[str, List[Tuple[str, str]]]:
                if not text:
                    return "", history
                try:
                    _, updated = respond(text, history, selected_topic, memory_flag)
                    return "", updated
                except Exception as e:
                    return "", history + [(text, f"Error: {str(e)}")]

            for btn in starter_buttons:
                btn.click(fn=use_starter, inputs=[btn, chatbot, topic, memory_flag], outputs=[msg, chatbot], queue=True)

            topic.change(fn=update_starters, inputs=[topic, difficulty], outputs=starter_buttons)
            difficulty.change(fn=update_starters, inputs=[topic, difficulty], outputs=starter_buttons)

    msg.submit(fn=respond, inputs=[msg, chatbot, topic, memory_flag], outputs=[msg, chatbot])
    submit.click(fn=respond, inputs=[msg, chatbot, topic, memory_flag], outputs=[msg, chatbot])
    clear.click(lambda: [], None, chatbot, queue=False)
    clear.click(lambda: "", None, msg, queue=False)

    default_starters = get_examples_for_topic("Daily Life", "intermediate")
    demo.load(fn=lambda: tuple(default_starters[:3]), outputs=starter_buttons, queue=False)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)