import gradio as gr
from openai import OpenAI
import os
import random
import asyncio

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

DEFAULT_LANGUAGE = "English"
WELCOME_MESSAGE_EN = "Hi! I'm your friendly assistant 🤖 Let's begin!\n\nPlease tell me which language you'd like to use."

# === Question Bank ===
SET_I = [
    "Given the choice of anyone in the world, whom would you want as a dinner guest?",
    "Would you like to be famous? In what way?",
    "Before making a telephone call, do you ever rehearse what you are going to say? Why?",
    "What would constitute a 'perfect' day for you?",
    "When did you last sing to yourself? To someone else?",
    "If you were able to live to the age of 90 and retain either the mind or body of a 30-year-old for the last 60 years of your life, which would you want?",
    "Do you have a secret hunch about how you will die?",
    "Name three things you and me appear to have in common.",
    "For what in your life do you feel most grateful?",
    "If you could change anything about the way you were raised, what would it be?",
    "Take one minutes and tell me your life story in as much detail as possible.",
    "If you could wake up tomorrow having gained any one quality or ability, what would it be?"
]

SET_II = [
    "If a crystal ball could tell you the truth about yourself, your life, the future or anything else, what would you want to know?",
    "Is there something that you’ve dreamed of doing for a long time? Why haven’t you done it?",
    "What is the greatest accomplishment of your life?",
    "What do you value most in a friendship?",
    "What is your most treasured memory?",
    "What is your most terrible memory?",
    "If you knew that in one year you would die suddenly, would you change anything about the way you are now living? Why?",
    "What does friendship mean to you?",
    "What roles do love and affection play in your life?",
    "Alternate sharing something you consider a positive characteristic of me. Share a total of three items.",
    "How close and warm is your family? Do you feel your childhood was happier than most other people’s?",
    "How do you feel about your relationship with your mother?"
]

SET_III = [
    "Make three true 'we' statements each. For instance, 'We are both in this chatroom feeling ...'",
    "Complete this sentence: 'I wish I had someone with whom I could share ...'",
    "If you were going to become a close friend with me, please share what would be important for me to know.",
    "Tell me what you like about me; be very honest this time, saying things that you might not say to someone you’ve just met.",
    "Share with me an embarrassing moment in your life.",
    "When did you last cry in front of another person? By yourself?",
    "Tell me something that you like about me already.",
    "What, if anything, is too serious to be joked about?",
    "If you were to die this evening with no opportunity to communicate with anyone, what would you most regret not having told someone? Why haven’t you told them yet?",
    "Your house, containing everything you own, catches fire. After saving your loved ones and pets, you have time to safely make a final dash to save any one item. What would it be? Why?",
    "Of all the people in your family, whose death would you find most disturbing? Why?",
    "Share a personal problem and ask my advice on how I might handle it. Also, ask me to reflect back to you how you seem to be feeling about the problem you have chosen."
]

def generate_question_set():
    return random.sample(SET_I, 3) + random.sample(SET_II, 3) + random.sample(SET_III, 3)

async def gpt_translate(text, target_lang):
    if target_lang == "English":
        return text
    prompt = f"Translate the following into {target_lang}:\n\n{text}"
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3
    )
    return response.choices[0].message.content.strip()

async def respond(user_input, history, step, language, questions, followup_mode):
    history.append({"role": "user", "content": user_input})
    if language == "":
        user_language = user_input.strip().capitalize()
        questions_en = generate_question_set()
        translated_questions = [await gpt_translate(q, user_language) for q in questions_en]
        welcome = await gpt_translate("Great! We will now continue in your selected language.", user_language)
        history.append({"role": "assistant", "content": f"{welcome}\n\n1. {translated_questions[0]}"})
        return history, "", 1, user_language, translated_questions, False
    if step >= len(questions):
        thank_you = await gpt_translate("Thank you for your response!", language)
        end_note = await gpt_translate("This concludes our questions. Please proceed with the rest of the survey. 📝", language)
        history.append({"role": "assistant", "content": thank_you})
        history.append({"role": "assistant", "content": end_note})
        return history, "", step, language, questions, False
    if followup_mode:
        comment_prompt = [
            {"role": "system", "content": f"Write a short, empathetic comment in {language} responding to the user's last answer."},
            history[-2],
            history[-1]
        ]
        comment_response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=comment_prompt,
            temperature=0.7
        )
        comment = comment_response.choices[0].message.content.strip()
        history.append({"role": "assistant", "content": ""})
        for c in comment:
            history[-1]["content"] += c
            await asyncio.sleep(0.03)
        next_question = f"{step+1}. {questions[step]}"
        history.append({"role": "assistant", "content": ""})
        for c in next_question:
            history[-1]["content"] += c
            await asyncio.sleep(0.03)
        return history, "", step + 1, language, questions, False

    comment_prompt = [
        {"role": "system", "content": f"You are a friendly and witty assistant. Write a short, personalized comment on the user's answer in {language}. Add a light emoji at the end."},
        history[-2],
        history[-1]
    ]
    comment_response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=comment_prompt,
        temperature=0.8
    )
    bot_reply = comment_response.choices[0].message.content.strip()

    followup_count = 0
    rand = random.random()
    if rand < 0.1:
        followup_count = 2
    elif rand < 0.5:
        followup_count = 1

    if followup_count > 0:
        followup_prompt = [
            {"role": "system", "content": f"Ask {followup_count} open-ended follow-up question(s) based on the user's last answer in {language}."},
            history[-2],
            history[-1]
        ]
        followup_response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=followup_prompt,
            temperature=0.7
        )
        bot_reply += "\n\n" + followup_response.choices[0].message.content.strip()
        history.append({"role": "assistant", "content": ""})
        for c in bot_reply:
            history[-1]["content"] += c
            await asyncio.sleep(0.03)
        return history, "", step, language, questions, True

    history.append({"role": "assistant", "content": ""})
    for c in bot_reply:
        history[-1]["content"] += c
        await asyncio.sleep(0.03)
    next_question = f"{step+1}. {questions[step]}"
    history.append({"role": "assistant", "content": ""})
    for c in next_question:
        history[-1]["content"] += c
        await asyncio.sleep(0.03)
    return history, "", step + 1, language, questions, False

def init():
    return [{"role": "assistant", "content": WELCOME_MESSAGE_EN}], 0, "", [], False

with gr.Blocks(css="""
.send-btn {
  background-color: #25D366 !important;
  color: white !important;
  border: none;
  border-radius: 24px;
  font-size: 20px;
  padding: 8px 20px;
  height: 48px;
  width: 60px;
  margin-left: 8px;
  cursor: pointer;
}
#chatbox .avatar-container {
    width: 64px !important;
    height: 64px !important;
    min-width: 64px !important;
    min-height: 64px !important;
    background-size: 100% 100% !important;
    background-position: center center !important;
    border-radius: 50% !important;
    padding: 0 !important;
    margin: 0 !important;
    border: none !important;
    box-shadow: none !important;
    background-color: transparent !important;
}
#chatbox .message.user { background-color: #DCF8C6 !important; }
#chatbox .message.assistant { background-color: #ffffff !important; }

""") as demo:
    gr.HTML("""
    <script>
      const waitForChatbox = () => {
        const chatbox = document.getElementById("chatbox");
        if (chatbox) {
          const observer = new MutationObserver(() => {
            chatbox.scrollTop = chatbox.scrollHeight;
          });
          observer.observe(chatbox, { childList: true, subtree: true });
        } else {
          setTimeout(waitForChatbox, 500);
        }
      };
      waitForChatbox();
    </script>
    """)

    chatbot = gr.Chatbot(
        elem_id="chatbox",
        label="",
        avatar_images=["user_avatar.jpg", "bot_avatar.png"],
        bubble_full_width=False,
        height=500,
        show_copy_button=False,
        type="messages"
    )

    with gr.Row():
        user_input = gr.Textbox(
            show_label=False,
            placeholder="Type your message here and press send...",
            container=True,
            scale=10
        )
        send_btn = gr.Button(value="➤", elem_classes="send-btn")

    state = gr.State([])
    step = gr.State(0)
    language = gr.State("")
    questions = gr.State([])
    followup_mode = gr.State(False)

    demo.load(fn=init, outputs=[chatbot, step, language, questions, followup_mode])
    user_input.submit(respond, [user_input, state, step, language, questions, followup_mode],
                      [chatbot, user_input, step, language, questions, followup_mode])
    send_btn.click(respond, [user_input, state, step, language, questions, followup_mode],
                   [chatbot, user_input, step, language, questions, followup_mode])

    demo.queue().launch()