import os
hf_token = os.getenv("Gem") # Store API token in .env
import streamlit as st
import os
import logging
import dotenv
import yaml
import PyPDF2
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFaceHub
import random
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

dotenv.load_dotenv()

# Load configuration from YAML
def load_config():
    with open("config.yaml", "r") as f:
        return yaml.safe_load(f)

config = load_config()
logging.basicConfig(level=logging.INFO)

# Load embedding model
embeddings_model = HuggingFaceEmbeddings(model_name=config["embedding_model"])

# Extract text from PDFs
def extract_text_from_pdf(file):
    reader = PyPDF2.PdfReader(file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() or ""
    return text.strip()

# Function to calculate matching score between job description and resume
def calculate_matching_score(jd_text, resume_text):
    vectorizer = TfidfVectorizer().fit_transform([jd_text, resume_text])
    score = cosine_similarity(vectorizer[0], vectorizer[1])[0][0] * 100
    return round(score, 2)

# Function to generate final score based on user responses
def calculate_final_score(responses):
    total_questions = len(responses)
    correct_responses = sum(1 for response in responses if "good" in response.lower() or "correct" in response.lower())
    return round((correct_responses / total_questions) * 100, 2) if total_questions > 0 else 0

# Get interview questions and assess responses
def get_interview_response(jd_text, resume_text, candidate_response=None, round_stage="intro", question_count=0):
    technical_names = ["Alex", "Jordan", "Casey", "Morgan"]
    hr_names = ["Taylor", "Jamie", "Riley", "Sam"]

    if round_stage in ["technical", "coding"]:
        interviewer_name = random.choice(technical_names)
        role = "Technical Lead"
    else:
        interviewer_name = random.choice(hr_names)
        role = "HR Manager"

    prompt_template = f"""
    My name is {interviewer_name}, and I am your {role} for this round.

    JOB DESCRIPTION:
    {jd_text}

    CANDIDATE PROFILE:
    {resume_text}

    This is question {question_count+1} of 5.

    """

    if question_count >= 5:
        return f"{interviewer_name}: This round is complete. Moving to the next stage."

    if round_stage == "intro":
        prompt_template += f"{interviewer_name}: Let's start with an introduction. Tell me about yourself."

    elif round_stage == "technical":
        prompt_template += f"{interviewer_name}: Based on your resume and the job description, here is a technical question for you."

    elif round_stage == "coding":
        prompt_template += f"{interviewer_name}: Let's move to a coding problem relevant to your role."

    elif round_stage == "hr":
        prompt_template += f"{interviewer_name}: Now let's discuss some HR aspects, starting with your motivation for this role."

    elif round_stage == "final_feedback":
        prompt_template += "Summarize the candidate’s performance in both rounds in a structured format."

    if candidate_response:
        if candidate_response.lower() == "hint":
            prompt_template += f"{interviewer_name}: Here is a helpful hint."
        else:
            prompt_template += f"The candidate answered: {candidate_response}. Assess the response and move to the next question."

    llm = HuggingFaceHub(
        repo_id=config["model_name"],
        model_kwargs={"temperature": config["temperature"], "max_length": 200},
        huggingfacehub_api_token=hf_token
    )

    response = llm(prompt_template).strip()

    # Store the full assessment in a text file for admin review
    with open("candidate_assessment.txt", "a") as f:
        f.write(f"Round: {round_stage}, Question {question_count+1}\n")
        f.write(f"Interviewer: {interviewer_name} ({role})\n")
        f.write(f"Question: {prompt_template}\n")
        f.write(f"Candidate Response: {candidate_response}\n")
        f.write(f"Feedback: {response}\n\n")

    return response if round_stage != "final_feedback" else f"{interviewer_name}: The interview is now complete."

# Streamlit UI
st.set_page_config(page_title="AI Interviewer", layout="centered")

st.title("🤖 AI Interview Chatbot")
st.write("Upload a Job Description and Resume to start the interview.")

jd_file = st.file_uploader("Upload Job Description (PDF)", type=["pdf"])
resume_file = st.file_uploader("Upload Candidate Resume (PDF)", type=["pdf"])

if jd_file and resume_file:
    jd_text = extract_text_from_pdf(jd_file)
    resume_text = extract_text_from_pdf(resume_file)

    # Calculate matching score
    matching_score = calculate_matching_score(jd_text, resume_text)

    # Store interview history & matching score
    if "interview_history" not in st.session_state:
        st.session_state["interview_history"] = []
        st.session_state["responses"] = []
        first_question = get_interview_response(jd_text, resume_text)
        st.session_state["interview_history"].append(("AI", first_question))

    st.write(f"**Matching Score:** {matching_score}%")

    for role, msg in st.session_state["interview_history"]:
        st.chat_message(role).write(msg)

    query = st.chat_input("Your Response:")

    if query:
        response = get_interview_response(jd_text, resume_text, query)
        st.session_state["interview_history"].append(("You", query))
        st.session_state["interview_history"].append(("AI", response))
        st.session_state["responses"].append(response)  # Store responses for final score
        st.rerun()

# Generate final score and store the results for download
if "responses" in st.session_state and len(st.session_state["responses"]) >= 5:
    final_score = calculate_final_score(st.session_state["responses"])

    # Store all results in a text file
    file_path = "candidate_assessment.txt"
    with open(file_path, "w") as f:
        f.write(f"Matching Score: {matching_score}%\n")
        f.write(f"Final Score: {final_score}%\n\n")
        f.write("Interview Assessment:\n")
        for role, msg in st.session_state["interview_history"]:
            f.write(f"{role}: {msg}\n")

    # Provide file download option
    with open(file_path, "rb") as f:
        st.download_button("Download Assessment", f, file_name="candidate_assessment.txt")