import streamlit as st
from openai import OpenAI
import os
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import torch
import requests

# Set up OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Set up ElevenLabs API key
ELEVENLABS_API_KEY = "your_api_key"

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load metadata and embeddings (ensure these files are in your working directory or update paths)
metadata_path = 'question_metadata.csv'  # Update this path if needed
embeddings_path = 'question_dataset_embeddings.npy'  # Update this path if needed

metadata = pd.read_csv(metadata_path)
embeddings = np.load(embeddings_path)

# Load the SentenceTransformer model
model = SentenceTransformer("all-MiniLM-L6-v2").to(device)

# Load prompts from files
with open("question_generation_prompt.txt", "r") as file:
    question_generation_prompt = file.read()

with open("technical_interviewer_prompt.txt", "r") as file:
    technical_interviewer_prompt = file.read()

st.title("Real-World Programming Question Mock Interview")

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

if "follow_up_mode" not in st.session_state:
    st.session_state.follow_up_mode = False  # Tracks whether we're in follow-up mode

if "generated_question" not in st.session_state:
    st.session_state.generated_question = None  # Stores the generated question for persistence

if "debug_logs" not in st.session_state:
    st.session_state.debug_logs = []  # Stores debug logs for toggling

if "code_output" not in st.session_state:
    st.session_state.code_output = None  # Stores the output of executed Python code

# Function to find the top 1 most similar question based on user input
def find_top_question(query):
    query_embedding = model.encode(query, convert_to_tensor=True, device=device).cpu().numpy()
    query_embedding = query_embedding.reshape(1, -1)  # Reshape to (1, n_features)
    similarities = cosine_similarity(query_embedding, embeddings).flatten()
    top_index = similarities.argsort()[-1]
    top_result = metadata.iloc[top_index].copy()
    top_result['similarity_score'] = similarities[top_index]
    return top_result

# Function to generate response using OpenAI API with debugging logs
def generate_response(messages):
    debug_log_entry = {"messages": messages}
    st.session_state.debug_logs.append(debug_log_entry)  # Store debug log
    
    response = client.chat.completions.create(
        model="o1-mini",
        messages=messages,
    )
    
    return response.choices[0].message.content

# Function to generate audio using ElevenLabs API
def generate_audio(text):
    url = "https://api.elevenlabs.io/v1/text-to-speech"
    headers = {
        "xi-api-key": ELEVENLABS_API_KEY,
        "content-type": "application/json"
    }
    payload = {
        "text": text,
        "voice_id": "21m00tcm4tlvdq8ikwam",  # Default voice ID; replace with desired voice ID.
        "voice_settings": {
            "similarity_boost": 0.85,
            "stability": 0.5
        }
    }
    
    response = requests.post(url, headers=headers, json=payload)
    
    if response.status_code == 200:
        audio_file_path = f"assistant_response.mp3"
        with open(audio_file_path, "wb") as audio_file:
            audio_file.write(response.content)
        return audio_file_path
    else:
        st.error(f"Error generating audio: {response.status_code} - {response.text}")
        return None

# User input form for generating a new question
with st.form(key="input_form"):
    company = st.text_input("Company", value="Google")
    difficulty = st.selectbox("Difficulty", ["Easy", "Medium", "Hard"], index=1)
    topic = st.text_input("Topic (e.g., Backtracking)", value="Backtracking")
    
    generate_button = st.form_submit_button(label="Generate")

if generate_button:
    st.session_state.messages = []
    st.session_state.follow_up_mode = False
    
    query = f"{company} {difficulty} {topic}"
    top_question = find_top_question(query)
    
    detailed_prompt = (
        f"Transform this LeetCode question into a real-world interview scenario:\n\n"
        f"**Company**: {top_question['company']}\n"
        f"**Question Name**: {top_question['questionName']}\n"
        f"**Difficulty Level**: {top_question['difficulty level']}\n"
        f"**Tags**: {top_question['Tags']}\n"
        f"**Content**: {top_question['Content']}\n"
        f"\nPlease create a real-world interview question based on this information."
    )
    
    response_text = generate_response([{"role": "assistant", "content": question_generation_prompt}, {"role": "user", "content": detailed_prompt}])
    
    st.session_state.generated_question = response_text

    st.session_state.messages.append({"role": "assistant", "content": response_text})
    
    st.session_state.follow_up_mode = True

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if st.session_state.follow_up_mode:
    if user_input := st.chat_input("Continue your conversation or ask follow-up questions here:"):
        with st.chat_message("user"):
            st.markdown(user_input)
        
        st.session_state.messages.append({"role": "user", "content": user_input})

        assistant_response_text = generate_response(
            [{"role": "assistant", "content": technical_interviewer_prompt}] + st.session_state.messages
        )

        assistant_audio_path = generate_audio(assistant_response_text)

        with st.chat_message("assistant"):
            st.markdown(assistant_response_text)
            if assistant_audio_path:
                audio_bytes = open(assistant_audio_path, "rb").read()
                st.audio(audio_bytes, format="audio/mp3")
        
        st.session_state.messages.append({"role": "assistant", "content": assistant_response_text})

# Left Sidebar: Generated Question and Code Box
with st.sidebar:
    # Top Half: Generated Question
    st.markdown("## Generated Question")
    if st.session_state.generated_question:
        st.markdown(st.session_state.generated_question)
    else:
        st.markdown("_No question generated yet._")

    # Divider between sections
    st.markdown("---")

    # Bottom Half: Python Code Box
    st.markdown("## Python Code Interpreter")
    
    code_input = st.text_area("Write your Python code here:")
    
    col1, col2 = st.columns(2)
    
    with col1:
        if st.button("Run Code"):
            try:
                exec_globals = {}
                exec(code_input, exec_globals)  # Execute user-provided code safely within its own scope.
                output_key_values = {k: v for k, v in exec_globals.items() if k != "__builtins__"}
                if output_key_values:
                    output_strs = [f"{key}: {value}" for key, value in output_key_values.items()]
                    output_display_strs = "\n".join(output_strs)
                    output_display_strs += "\nCode executed successfully!"
                    print(output_display_strs)

            except Exception as e: