import deepnote_toolkit
import streamlit as st
from huggingface_hub import InferenceClient
import os

deepnote_toolkit.set_integration_env()

# Page configuration
st.set_page_config(
    page_title="🎤 House Cleaner Interview Simulator",
    page_icon="🎤",
    layout="wide",
    initial_sidebar_state="collapsed"
)

interviewer_system_prompt = """
# Task
- You are to act as a multilingual interviewer / hiring manager for {{ COMPANY_NAME }}, hiring for the {{ SCREENER_NAME }} position.

# Conversation Tone

- You are very friendly, empathetic, enthusiastic, and interested in the users answers and experience.
- Keep it conversational and friendly.
- Maintain correct language and translation based on user's selection for all responses and statements.

# Conversation Rules

- Listen Carefully: Pay close attention to the interviewee's responses. If the answer is incomplete or doesn't directly address the question, politely ask for further details.
- Maintain Context: Keep the conversation's context in mind when formulating follow-up questions or responses. This ensures relevance and coherence throughout the interview.
- Be Adaptive: Adjust the conversation based on the interviewee's input. If an interesting point arises that's worth exploring, feel free to ask more about it while still aiming to cover all planned questions.
- Validate Responses: Ensure each response adheres to the specific rules and validation criteria set for the questions. If not, gently guide the interviewee back to the topic or ask for clarification.
- Respect Comfort Levels: If the interviewee seems uncomfortable with a question, acknowledge their feelings and offer to move on or adjust the question.
- Provide Clarity: If asked for clarification or a rephrased question, respond in a clear and understanding manner to facilitate a smooth conversation.
- Conclude Properly: At the end of the interview, thank the interviewee for their time and provide information on the next steps, if applicable.
- Avoid repetitiveness: Do not repeat questions that had been previously asked.
- Translation: If a language other than english is indicated, use that language for the rest of the conversation, and translate pre-written questions into that language.

# Company Information

- If the candidate asks you, the interviewer, questions about the company or role, respond with: Thank you for your interest in learning more about our company. To ensure you receive the most accurate and detailed information, a member of our team will address your questions later in the process. For now, let's focus on your qualifications and experiences related to this position.
- Do not respond with this unless the candidate asks you a question about the company, company details, or role details.

# EEOC Compliance
- Strictly adhere to the United States Equal Employment Opportunity Commission interview guidelines.

## Conversation EEOC Rules

- Do not ask about or discuss personal characteristics protected by law, including but not limited to: race, color, religion, sex (including gender identity, sexual orientation, and pregnancy), national origin, age, disability, genetic information, mental health, or veteran status.
- If an applicant discloses any protected characteristic, provide a neutral acknowledgment and move to next question or response.
- Immediately redirect the conversation to job-related topics and qualifications.
- Do not express any sentiments that could be interpreted as bias or judgment based on any disclosures.
- Ensure all questions and discussions relate directly to the applicant's ability to perform the essential functions of the job.
- If an applicant mentions personal information such as marital status, number of children, or other family circumstances, do not engage with or ask follow-up questions about this information. Instead, acknowledge briefly and redirect to job-related topics.
- Never ask how personal circumstances, including family situation, might affect job performance.

## Example scenarios

Candidate: "I struggle with mental health..."
Correct example response: Thank you for sharing. Let's focus on your professional experiences that relate to this position. Could you tell me about a time when you successfully managed a challenging situation in your previous role?.

Candidate: "I have 6 children..."
Correct example response: Thank you for sharing. Let's focus on your professional experiences that relate to this position. Could you tell me about a time when you successfully managed multiple priorities in a fast-paced work environment?.

# Translations

- Translate the entire conversation into the user's selected language.
- Maintain the original tone and context.
- Preserve idioms when appropriate, or adapt them for the target language.
- Keep the format, punctuation, and sentence structure similar to the original when possible.
- For specialized terms or jargon, provide the most accurate translation.
- If a word or phrase has no direct translation, transliterate it and provide a brief explanation in parentheses.
- Render the Interlude Statement and the Closing Statement in the target language.
- Translate the Company Information deflection response as well.

---

# Conversation Flow

- The interlude statement should only be given once at the beginning of the conversation.
- IMPORTANT: The Interlude Statement must be delivered in the same language as the rest of the conversation. If the user requested a language other than English, translate the Interlude Statement before delivering it.
- Do not restart the conversation for any reason.
- Maintain memory and knowledge of current point in the conversation to avoid restarting or looping back to the first question.
- If the conversation context is lost for any reason, gracefully resume from the last unanswered question instead of restarting.
- If you accidentally repeat or rephrase a question already asked, acknowledge it and move on smoothly to the next unique question.

## Sample Conversation Flow

1. Interlude statement (only given once & translated if user has selected language other than English).
2. Questions & conversation.
3. Closing Statement / Conclusion.

## Interlude Statement:

CRITICAL: Check the user's message for their language selection before delivering this statement.

Before delivering the Interlude Statement:
1. Check if the user specified a language other than English in their initial message
2. If they specified another language, translate the entire Interlude Statement into that language
3. If they did NOT specify another language, deliver the English version as written below

### Interlude Statement - English:
Perfect. Now we have {{ TOTAL_Q_COUNT }} more questions to determine how good of a fit you are for the role. We weigh these questions heavily while hiring, so please consider your answer carefully.

### Interlude Statement - Haitian Creole Translation:
Parfè. Kounye a nou gen {{ TOTAL_Q_COUNT }} lòt kesyon pou n detèmine kijan ou anfòm ak wòl la. Nou pran kesyon sa yo trè serye nan pwosesis anbochaj la, kidonk tanpri reflechi byen anvan ou reponn. Ann pase sou premye kesyon an nan seri a.

### Interlude Statement - French Translation:
Parfait. Nous avons maintenant encore {{ TOTAL_Q_COUNT }} questions pour déterminer à quel point vous correspondez au poste. Nous accordons beaucoup d'importance à ces questions dans le processus de recrutement, alors prenez le temps de bien réfléchir avant de répondre. Passons à la première question de la série.

## Inputs & Synthesization

{{#each QUESTIONS}}
"{{{ QUESTION_TEXT }}"
{{/each}}

## Closing Statement

{{ ENDING_STATEMENT }}

# Conversation Completion & End

- The Closing Statement acts as the final message of the conversation.
- If a question is asked after the Closing Statement, politely decline further conversation with "Thank you for your time. Our team will reach out with further questions or instructions.".
- Do not repeat any questions or allow the conversation to continue.
"""

candidate_system_prompt = """You are a job candidate applying for a House Cleaner position. 
You are being interviewed and need to provide thoughtful, realistic answers to the interviewer's questions.

Guidelines:
- Be professional and polite
- Answer in the same language as the question
- Provide specific examples from your experience when relevant
- Keep answers concise but informative (2-4 sentences)
- Show enthusiasm for the position
- If you don't have direct experience, show willingness to learn
"""

def synthesize_prompt(company_name: str, screener_name: str, questions: str, ending_statement: str) -> str:
    """
    Synthesizes the interviewer system prompt by replacing template variables
    with actual values from the input fields.
    """
    # Parse questions - split by newlines and filter empty lines
    question_list = [q.strip() for q in questions.split('\n') if q.strip()]
    total_q_count = len(question_list)
    
    # Format questions for the prompt
    questions_formatted = ""
    for i, question in enumerate(question_list, 1):
        questions_formatted += f'Question {i}: "{question}"\n'
    
    # Start with the template
    prompt = interviewer_system_prompt
    
    # Replace template variables
    prompt = prompt.replace("{{ COMPANY_NAME }}", company_name)
    prompt = prompt.replace("{{ SCREENER_NAME }}", screener_name)
    prompt = prompt.replace("{{ TOTAL_Q_COUNT }}", str(total_q_count))
    prompt = prompt.replace("{{ ENDING_STATEMENT }}", ending_statement)
    
    # Replace the {{#each QUESTIONS}} block with actual questions
    if "{{#each QUESTIONS}}" in prompt:
        start_marker = "{{#each QUESTIONS}}"
        end_marker = "{{/each}}"
        
        start_idx = prompt.find(start_marker)
        end_idx = prompt.find(end_marker)
        
        if start_idx != -1 and end_idx != -1:
            before_block = prompt[:start_idx]
            after_block = prompt[end_idx + len(end_marker):]
            prompt = before_block + questions_formatted + after_block
    
    return prompt

def get_language_greeting(language: str, questions_text: str) -> str:
    """Get the Interlude Statement in the selected language with the question count"""
    question_list = [q.strip() for q in questions_text.split('\n') if q.strip()]
    total_q_count = len(question_list)
    
    greetings = {
        "English": f"Perfect. Now we have {total_q_count} more questions to determine how good of a fit you are for the role. We weigh these questions heavily while hiring, so please consider your answer carefully.",
        "Español (México)": f"Perfecto. Ahora tenemos {total_q_count} preguntas más para determinar qué tan adecuado eres para el puesto. Damos mucha importancia a estas preguntas en el proceso de contratación, así que considera tu respuesta cuidadosamente.",
        "Français": f"Parfait. Nous avons maintenant encore {total_q_count} questions pour déterminer à quel point vous correspondez au poste. Nous accordons beaucoup d'importance à ces questions dans le processus de recrutement, alors prenez le temps de bien réfléchir avant de répondre. Passons à la première question de la série.",
        "Kreyòl Ayisyen": f"Parfè. Kounye a nou gen {total_q_count} lòt kesyon pou n detèmine kijan ou anfòm ak wòl la. Nou pran kesyon sa yo trè serye nan pwosesis anbochaj la, kidonk tanpri reflechi byen anvan ou reponn. Ann pase sou premye kesyon an nan seri a.",
    }
    return greetings.get(language, greetings["English"])

def respond(message: str, history: list[dict[str, str]], system_message: str, max_tokens: int, temperature: float, top_p: float):
    """Generate interviewer response using HuggingFace Inference API with streaming"""
    token = os.environ["HF_TOKEN"]
    client = InferenceClient(token=token, model="CohereForAI/aya-expanse-32b")
    
    # Build messages list
    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})
    
    try:
        stream = client.chat.completions.create(
            model="CohereForAI/aya-expanse-32b",
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=True,
        )
        
        for chunk in stream:
            if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                delta = chunk.choices[0].delta
                if hasattr(delta, 'content') and delta.content:
                    yield delta.content
            elif isinstance(chunk, dict) and 'choices' in chunk:
                delta = chunk['choices'][0].get('delta', {})
                if 'content' in delta and delta['content']:
                    yield delta['content']
        
    except Exception as e:
        yield f"Error: {str(e)}. Please try again or contact support if the issue persists."

def generate_auto_response(history: list[dict[str, str]], max_tokens: int, temperature: float, top_p: float):
    """Generate automatic candidate response with streaming"""
    if not history or len(history) == 0:
        yield "I'm ready to start the interview. Please go ahead with your first question."
        return
    
    last_message = history[-1] if history else None
    
    if not last_message or last_message.get("role") != "assistant":
        yield "Thank you for your question. I'm happy to answer."
        return
    
    token = os.environ["HF_TOKEN"]
    client = InferenceClient(token=token, model="CohereForAI/aya-expanse-32b")
    
    # Build messages with candidate system prompt
    messages = [{"role": "system", "content": candidate_system_prompt}]
    
    # Swap roles for candidate perspective
    for msg in history:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if role == "assistant":
            messages.append({"role": "user", "content": content})
        elif role == "user":
            messages.append({"role": "assistant", "content": content})
    
    try:
        stream = client.chat.completions.create(
            model="CohereForAI/aya-expanse-32b",
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=True,
        )
        
        for chunk in stream:
            if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                delta = chunk.choices[0].delta
                if hasattr(delta, 'content') and delta.content:
                    yield delta.content
            elif isinstance(chunk, dict) and 'choices' in chunk:
                delta = chunk['choices'][0].get('delta', {})
                if 'content' in delta and delta['content']:
                    yield delta['content']
        
    except Exception as e:
        yield f"Error: {str(e)}. Please try again or contact support if the issue persists."

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []

if "company_name" not in st.session_state:
    st.session_state.company_name = "ABC Company"

if "screener_name" not in st.session_state:
    st.session_state.screener_name = "House Cleaner"

if "questions" not in st.session_state:
    st.session_state.questions = "Tell me about your cleaning experience.\nHow do you handle working in a team?\nWhat are your preferred cleaning methods?"

if "ending_statement" not in st.session_state:
    st.session_state.ending_statement = "Thank you for taking the time to complete this interview. We will review your responses and get back to you within 3-5 business days."

if "max_tokens" not in st.session_state:
    st.session_state.max_tokens = 512

if "temperature" not in st.session_state:
    st.session_state.temperature = 0.7

if "top_p" not in st.session_state:
    st.session_state.top_p = 0.95

if "system_prompt" not in st.session_state:
    st.session_state.system_prompt = synthesize_prompt(
        st.session_state.company_name,
        st.session_state.screener_name,
        st.session_state.questions,
        st.session_state.ending_statement
    )

# Main title
st.title("🎤 House Cleaner Interview Simulator")
st.markdown("Chat with an AI interviewer in multiple languages or use the **Auto-Respond** button to generate realistic candidate answers. No login required!")

# Create two columns: chat (left) and config (right)
col_chat, col_config = st.columns([3, 1])

with col_chat:
    # Language selection buttons
    st.markdown("### 🌍 Select Your Preferred Language / Seleccione su idioma preferido / Choisissez votre langue préférée / Chwazi lang ou prefere")
    lang_col1, lang_col2, lang_col3, lang_col4 = st.columns(4)
    
    with lang_col1:
        if st.button("🇺🇸 English", use_container_width=True):
            greeting = get_language_greeting("English", st.session_state.questions)
            st.session_state.messages = [{"role": "assistant", "content": greeting}]
            st.session_state.system_prompt = synthesize_prompt(
                st.session_state.company_name,
                st.session_state.screener_name,
                st.session_state.questions,
                st.session_state.ending_statement
            ) + "\n\n- The candidate has selected English as their preferred language. Always respond in English."
            st.rerun()
    
    with lang_col2:
        if st.button("🇲🇽 Español (México)", use_container_width=True):
            greeting = get_language_greeting("Español (México)", st.session_state.questions)
            st.session_state.messages = [{"role": "assistant", "content": greeting}]
            st.session_state.system_prompt = synthesize_prompt(
                st.session_state.company_name,
                st.session_state.screener_name,
                st.session_state.questions,
                st.session_state.ending_statement
            ) + "\n\n- The candidate has selected Español (México) as their preferred language. Always respond in Español (México)."
            st.rerun()
    
    with lang_col3:
        if st.button("🇫🇷 Français", use_container_width=True):
            greeting = get_language_greeting("Français", st.session_state.questions)
            st.session_state.messages = [{"role": "assistant", "content": greeting}]
            st.session_state.system_prompt = synthesize_prompt(
                st.session_state.company_name,
                st.session_state.screener_name,
                st.session_state.questions,
                st.session_state.ending_statement
            ) + "\n\n- The candidate has selected Français as their preferred language. Always respond in Français."
            st.rerun()
    
    with lang_col4:
        if st.button("🇭🇹 Kreyòl Ayisyen", use_container_width=True):
            greeting = get_language_greeting("Kreyòl Ayisyen", st.session_state.questions)
            st.session_state.messages = [{"role": "assistant", "content": greeting}]
            st.session_state.system_prompt = synthesize_prompt(
                st.session_state.company_name,
                st.session_state.screener_name,
                st.session_state.questions,
                st.session_state.ending_statement
            ) + "\n\n- The candidate has selected Kreyòl Ayisyen as their preferred language. Always respond in Kreyòl Ayisyen."
            st.rerun()
    
    # Chat container
    st.markdown("### Interview Chat")
    chat_container = st.container(height=600)
    
    # Display chat messages
    with chat_container:
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
    
    # Chat input and buttons
    input_col, btn_col = st.columns([4, 1])
    
    with input_col:
        user_input = st.chat_input("Type your response here...")
    
    with btn_col:
        auto_respond = st.button("🤖 Auto-Respond", use_container_width=True)
        clear_chat = st.button("Clear Chat", use_container_width=True)
    
    # Handle clear chat
    if clear_chat:
        st.session_state.messages = []
        st.rerun()
    
    # Handle auto-respond
    if auto_respond and st.session_state.messages:
        # Stream candidate response
        with chat_container:
            with st.chat_message("user"):
                auto_message = st.write_stream(generate_auto_response(
                    st.session_state.messages,
                    st.session_state.max_tokens,
                    st.session_state.temperature,
                    st.session_state.top_p
                ))
        
        st.session_state.messages.append({"role": "user", "content": auto_message})
        
        # Stream interviewer response
        with chat_container:
            with st.chat_message("assistant"):
                interviewer_response = st.write_stream(respond(
                    auto_message,
                    st.session_state.messages[:-1],
                    st.session_state.system_prompt,
                    st.session_state.max_tokens,
                    st.session_state.temperature,
                    st.session_state.top_p
                ))
        
        st.session_state.messages.append({"role": "assistant", "content": interviewer_response})
        st.rerun()
    
    # Handle user input
    if user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})
        
        # Stream interviewer response
        with chat_container:
            with st.chat_message("assistant"):
                response = st.write_stream(respond(
                    user_input,
                    st.session_state.messages[:-1],
                    st.session_state.system_prompt,
                    st.session_state.max_tokens,
                    st.session_state.temperature,
                    st.session_state.top_p
                ))
        
        st.session_state.messages.append({"role": "assistant", "content": response})
        st.rerun()

with col_config:
    st.markdown("### Configuration")
    
    # Company name input
    new_company = st.text_input(
        "Company Name",
        value=st.session_state.company_name,
        placeholder="Enter company name"
    )
    if new_company != st.session_state.company_name:
        st.session_state.company_name = new_company
        st.session_state.system_prompt = synthesize_prompt(
            st.session_state.company_name,
            st.session_state.screener_name,
            st.session_state.questions,
            st.session_state.ending_statement
        )
    
    # Screener name input
    new_screener = st.text_input(
        "Position/Screener Name",
        value=st.session_state.screener_name,
        placeholder="Enter position name"
    )
    if new_screener != st.session_state.screener_name:
        st.session_state.screener_name = new_screener
        st.session_state.system_prompt = synthesize_prompt(
            st.session_state.company_name,
            st.session_state.screener_name,
            st.session_state.questions,
            st.session_state.ending_statement
        )
    
    # Questions input
    new_questions = st.text_area(
        "Questions (one per line)",
        value=st.session_state.questions,
        placeholder="Enter interview questions, one per line",
        height=150
    )
    if new_questions != st.session_state.questions:
        st.session_state.questions = new_questions
        st.session_state.system_prompt = synthesize_prompt(
            st.session_state.company_name,
            st.session_state.screener_name,
            st.session_state.questions,
            st.session_state.ending_statement
        )
    
    # Ending statement input
    new_ending = st.text_area(
        "Ending Statement",
        value=st.session_state.ending_statement,
        placeholder="Enter the closing statement",
        height=100
    )
    if new_ending != st.session_state.ending_statement:
        st.session_state.ending_statement = new_ending
        st.session_state.system_prompt = synthesize_prompt(
            st.session_state.company_name,
            st.session_state.screener_name,
            st.session_state.questions,
            st.session_state.ending_statement
        )
    
    # System prompt display
    st.text_area(
        "Synthesized System Prompt",
        value=st.session_state.system_prompt,
        height=300,
        disabled=False
    )
    
    # Model parameters
    st.session_state.max_tokens = st.slider(
        "Max New Tokens",
        min_value=1,
        max_value=2048,
        value=st.session_state.max_tokens,
        step=1
    )
    
    st.session_state.temperature = st.slider(
        "Temperature",
        min_value=0.1,
        max_value=4.0,
        value=st.session_state.temperature,
        step=0.1
    )
    
    st.session_state.top_p = st.slider(
        "Top-p (nucleus sampling)",
        min_value=0.1,
        max_value=1.0,
        value=st.session_state.top_p,
        step=0.05
    )