import json
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import os

# Set up device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

hf_token = os.environ["HF_TOKEN"]

# Load model and tokenizer from local files
model_path = "AI-Mock-Interviewer/T5"  # Assuming all files are in the root directory
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

# Move model to the appropriate device
model.to(device)

# ------------------- SYSTEM PROMPT -------------------
system_prompt = """
You are conducting a mock technical interview. The candidate's experience level can be entry-level, mid-level, or senior-level. Generate questions and follow-up questions based on the domain and the candidate's experience level. Consider these aspects:
1. The question should be relevant to the domain (e.g., software engineering, machine learning) and appropriate for the candidate's experience level.
2. For follow-up questions, analyze the candidate's last response and ask questions that probe deeper into their understanding, challenge their approach, or request clarification.
3. The follow-up question should aim to explore the candidate's depth of knowledge and ability to adapt.
4. Ensure each question is unique and does not repeat previously asked questions.
5. Ensure each question covers a different sub-topic within the domain, avoiding redundancy.
6. If no clear follow-up can be derived, generate a fresh, related question from a different aspect of the domain.
Important: Ensure that each question is clear, concise, and allows the candidate to demonstrate their technical and communicative abilities effectively.
"""

# Define sub-topic categories for different domains
subtopic_keywords = {
    "data analysis": ["data cleaning", "missing data", "EDA", "visualization"],
    "machine learning": ["supervised learning", "overfitting", "hyperparameter tuning"],
    "software engineering": ["code optimization", "design patterns", "database design"],
}

def identify_subtopic(question, domain):
    """Identify the sub-topic of a question using predefined keywords."""
    domain = domain.lower()
    if domain in subtopic_keywords:
        for subtopic in subtopic_keywords[domain]:
            if subtopic in question.lower():
                return subtopic
    return None

# Tracking asked questions
def generate_question(prompt, domain, state=None):
    """
    Generates a unique question based on the prompt and domain.
    Uses 'state' to track uniqueness in the conversation session.
    """
    full_prompt = system_prompt + "\n" + prompt
    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
    
    outputs = model.generate(
        inputs["input_ids"],
        max_new_tokens=50,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        top_k=30,
        top_p=0.9,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )
    
    question = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    # Ensure question ends with a question mark
    if not question.endswith("?"):
        question += "?"

    # Identify the subtopic
    subtopic = identify_subtopic(question, domain)

    # Check for uniqueness
    if state is not None:
        if (question not in state["asked_questions"] and
            (subtopic is None or subtopic not in state["asked_subtopics"])):
            state["asked_questions"].add(question)
            if subtopic:
                state["asked_subtopics"].add(subtopic)
            return question
    return question  # Fallback

# Initialize conversation state
def reset_state(domain, company, level):
    return {
        "domain": domain,
        "company": company,
        "level": level,
        "asked_questions": set(),
        "asked_subtopics": set(),
        "conversation": []  # List of (speaker, message) tuples
    }

def start_interview(domain, company, level):
    state = reset_state(domain, company, level)
    prompt = f"Domain: {domain}. Candidate experience level: {level}. Generate the first question:"
    
    question = generate_question(prompt, domain, state)
    state["conversation"].append(("Interviewer", question))
    return state["conversation"], state

def submit_response(candidate_response, state):
    state["conversation"].append(("Candidate", candidate_response))
    prompt = (f"Domain: {state['domain']}. Candidate's last response: {candidate_response}. Generate a follow-up question:")
    
    question = generate_question(prompt, state["domain"], state)
    state["conversation"].append(("Interviewer", question))
    return state["conversation"], state

# ----------- Gradio UI -----------

with gr.Blocks() as demo:
    gr.Markdown("# Interactive AI-Powered Mock Interview")
    
    with gr.Row():
        domain_input = gr.Textbox(label="Domain", placeholder="e.g. Software Engineering")
        company_input = gr.Textbox(label="Company (Optional)", placeholder="e.g. Google")
        level_input = gr.Dropdown(
            label="Experience Level",
            choices=["Entry-Level", "Mid-Level", "Senior-Level"],
            value="Entry-Level"
        )

    start_button = gr.Button("Start Interview")
    chatbot = gr.Chatbot(label="Interview Conversation")

    with gr.Row():
        response_input = gr.Textbox(label="Your Response")
        submit_button = gr.Button("Submit")
        clear_button = gr.Button("Clear Chat")

    state = gr.State()  # Holds session data

    start_button.click(start_interview, inputs=[domain_input, company_input, level_input], outputs=[chatbot, state])
    submit_button.click(submit_response, inputs=[response_input, state], outputs=[chatbot, state]).then(lambda: "", None, response_input)
    clear_button.click(lambda: ([], None), outputs=[chatbot, state])

demo.launch()