import os
import streamlit as st
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
from sentence_transformers import SentenceTransformer, util

# --- CONFIGURATION ---
NUM_QUESTIONS_TO_ASK = 5 # <--- NEW: Set the limit here

# --- HF cache setup in a writable folder (/tmp is safe in Docker) ---
HF_CACHE = "/tmp/hf_cache"
os.makedirs(HF_CACHE, exist_ok=True)
os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
os.environ["HF_HOME"] = HF_CACHE
os.environ["HF_DATASETS_CACHE"] = HF_CACHE
os.environ["HF_MODULES_CACHE"] = HF_CACHE

st.set_page_config(page_title="Interviewer Bot", layout="wide")
st.title("📝 Interviewer Bot")
st.write(f"AI-powered Interviewer Bot for {NUM_QUESTIONS_TO_ASK} random questions.")

# --- Load fine-tuned model (Explicitly passing cache_dir & fixing FutureWarnings) ---
@st.cache_resource(show_spinner=True)
def load_model(model_name, cache_dir):
    # FIX: Explicitly setting clean_up_tokenization_spaces to suppress FutureWarning
    tokenizer = AutoTokenizer.from_pretrained(
        model_name, 
        cache_dir=cache_dir,
        clean_up_tokenization_spaces=True
    )
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, cache_dir=cache_dir)
    return tokenizer, model

MODEL_NAME = "aswinprasath31/interviewer-bot-finetuned-v1"
tokenizer, model = load_model(MODEL_NAME, HF_CACHE)

# --- Load and preprocess questions from a CSV file (Data Fix for combined Q&A) ---
@st.cache_data
def load_questions(csv_path="training_data.csv"):
    df = pd.read_csv(csv_path, header=None) 
    
    if df.iloc[0, 0] == 'test':
        df = df.iloc[1:].copy()

    df[['Question', 'Answer']] = df[0].str.split('Answer:', expand=True)
    df['Question'] = df['Question'].str.replace('Question:', '').str.strip()
    df['Answer'] = df['Answer'].str.strip()
    df.dropna(subset=['Question'], inplace=True)
    
    return df[['Question', 'Answer']].reset_index(drop=True)

# Load the full question set once
full_questions_df = load_questions()

# --- Session state management and Random Sampling (CRITICAL CHANGE) ---
if "sampled_questions" not in st.session_state:
    # Randomly sample the specified number of questions (without replacement)
    # The 'min' ensures we don't crash if the CSV has fewer than 5 questions
    N = min(NUM_QUESTIONS_TO_ASK, len(full_questions_df))
    st.session_state.sampled_questions = full_questions_df.sample(n=N).reset_index(drop=True)

if "q_index" not in st.session_state:
    st.session_state.q_index = 0
if "responses" not in st.session_state:
    st.session_state.responses = []

# Now, all subsequent logic uses the sampled_questions dataframe
questions_df = st.session_state.sampled_questions
num_questions = len(questions_df)


# --- Sentence transformer for similarity scoring (Explicitly passing cache_dir) ---
@st.cache_resource(show_spinner=False)
def get_embedder(cache_folder):
    return SentenceTransformer('all-MiniLM-L6-v2', cache_folder=cache_folder)

embedder = get_embedder(HF_CACHE)

def get_reference_answer(question, max_len=200):
    inputs = tokenizer(question, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=max_len)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def compute_similarity(candidate_answer, reference_answer):
    candidate_emb = embedder.encode(candidate_answer, convert_to_tensor=True)
    reference_emb = embedder.encode(reference_answer, convert_to_tensor=True)
    score = util.cos_sim(candidate_emb, reference_emb).item()
    return round(score * 10, 2)  # scale 0-10

# --- UI Layout ---
if st.session_state.q_index < num_questions:
    st.progress(st.session_state.q_index / num_questions)
    # Access the current question from the sampled list
    question = questions_df.iloc[st.session_state.q_index]["Question"] 

    st.subheader(f"Question {st.session_state.q_index + 1}/{num_questions}")
    st.write(f"💡 {question}")

    candidate_answer = st.text_area("Your Answer:")

    if st.button("Submit Answer"):
        if candidate_answer.strip() == "":
            st.warning("⚠️ Please enter an answer before submitting.")
        else:
            # Generate the reference answer using your fine-tuned LLM
            reference_answer = get_reference_answer(question) 
            
            score = compute_similarity(candidate_answer, reference_answer)
            
            st.success("✅ Reference Answer:")
            st.write(reference_answer)
            st.info(f"📝 Score (0-10) based on similarity: {score}")
            
            # Save response
            st.session_state.responses.append({
                "Question": question,
                "Candidate Answer": candidate_answer,
                "Reference Answer": reference_answer,
                "Score": score
            })
            
            # Move to the next question
            st.session_state.q_index += 1
            st.rerun() # Rerun to display the next question or the completion message
            
else:
    # Completion message
    st.balloons()
    st.success("🎉 Interview completed!")
    
    # Display results and download button
    all_responses = pd.DataFrame(st.session_state.responses)
    avg_score = all_responses['Score'].mean()
    st.markdown(f"### Overall Average Score: **{avg_score:.2f}**")
    st.dataframe(all_responses, hide_index=True)
    
    csv = all_responses.to_csv(index=False).encode('utf-8')
    st.download_button(
        "📥 Download Responses CSV", 
        data=csv, 
        file_name="interview_responses.csv",
        type="primary"
    )

    # Optional: Reset button to start a new random interview
    if st.button("Start New Interview"):
        del st.session_state.sampled_questions
        del st.session_state.q_index
        del st.session_state.responses
        st.rerun()