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 Robin, your dialogue partner.\nWe're about to begin a conversation that includes nine questions. Are you ready?\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?",
    "Have you ever thought about how you would like to be remembered in the future?",
    "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?",
    "Could you please tell me one of your life stories in detail?",
    "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 you think?",
    "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 were going to Mars and wouldn’t be coming back, would you change the way you’re currently living? Why?",
    "What does friendship mean to you?",
    "What roles do love and affection play in your life?",
    "Please share what you think are my strengths so far, I would be grateful.",
    "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 = [
    "Let's try something fun: Can you write one sentence that start with 'We' and describe something true about us — you and me, chatting here together? (For example: 'We are all happy with this conversation.')",
    "Let’s complete this sentence together: 'I wish I had someone I could share ___ with.' What would you say?",
    "If you were going to become a close friend with me, please share what would be important for me to know.",
    "Tell me how you feel about me",
    "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 leaving for Mars tonight and wouldn’t be able to contact anyone again, 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?",
    "If one member of your family had to go on a long mission to Mars and you wouldn’t see them for many years, whose absence would affect you the most? Why?",
    "Think of a personal problem you're facing right now, big or small, and tell me about it. I'd love to hear how you're feeling about it — and if you'd like, I can share how I might deal with something similar."
]

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": ""})
        for c in f"{welcome}\n\n1. {translated_questions[0]}":
            history[-1]["content"] += c
            await asyncio.sleep(0.03)
        return history, "", 1, user_language, translated_questions, False

    if step >= len(questions):
        thank_you = await gpt_translate("Thank you for your response! It was a pleasure talking with you. I hope you enjoyed the conversation.", language)
        end_note = await gpt_translate("This concludes our questions. Please proceed with the rest of the survey. 📝", language)
        for msg in [thank_you, end_note]:
            history.append({"role": "assistant", "content": ""})
            for c in msg:
                history[-1]["content"] += c
                await asyncio.sleep(0.03)
        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()

        rand = random.random()
        if rand < 0.1:
            followup_prompt = [
                {"role": "system", "content": f"Ask one open-ended follow-up question in {language} based on the user's last answer."},
                history[-2],
                history[-1]
            ]
            followup_response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=followup_prompt,
                temperature=0.7
            )
            followup_question = followup_response.choices[0].message.content.strip()

            history.append({"role": "assistant", "content": ""})
            for c in comment:
                history[-1]["content"] += c
                await asyncio.sleep(0.03)

            history.append({"role": "assistant", "content": ""})
            for c in followup_question:
                history[-1]["content"] += c
                await asyncio.sleep(0.03)

            return history, "", step, language, questions, True
        else:
            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"Write a short, personalized comment in {language} responding to the user's answer. Add a light emoji."},
        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()

    rand = random.random()
    if rand < 0.4:
        followup_prompt = [
            {"role": "system", "content": f"Ask one open-ended follow-up question in {language} based on the user's answer."},
            history[-2],
            history[-1]
        ]
        followup_response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=followup_prompt,
            temperature=0.7
        )
        followup_question = followup_response.choices[0].message.content.strip()

        history.append({"role": "assistant", "content": ""})
        for c in comment:
            history[-1]["content"] += c
            await asyncio.sleep(0.03)

        history.append({"role": "assistant", "content": ""})
        for c in followup_question:
            history[-1]["content"] += c
            await asyncio.sleep(0.03)

        return history, "", step, language, questions, True
    else:
        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

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: 54px !important;
    height: 54px !important;
    min-width: 54px !important;
    min-height: 54px !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 {
  display: inline-block !important;
  font-size: 16px !important;
  line-height: 1.3em !important;
  min-width: 7ch !important;
  max-width: 60ch !important;
  padding: 2px 3px !important;
  border-radius: 8px !important;
  word-break: break-word !important;
  white-space: normal !important;
}
@media screen and (max-width: 768px) {
  #chatbox .message.user {
    max-width: 90vw !important;
  }
}
#chatbox .message.user {
  background-color: #DCF8C6 !important;
  min-width: 60px !important;         
  max-width: fit-content !important;  
}
#chatbox .message.assistant { background-color: #ffffff !important; }
footer { display: none !important; }
.svelte-1ipelgc { display: none !important; }
.icon-button-wrapper {
  display: none !important;
}
label.svelte-1to105q {
  display: none !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()