import gradio as gr
import pandas as pd
import numpy as np
import os
import re
import torch
import base64
from transformers import pipeline
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

print("Loading E5 Retrieval Model and Embeddings...")

base_dir = os.path.dirname(__file__)
csv_path = os.path.join(base_dir, 'interview_forge_v3_complete.csv')
if not os.path.exists(csv_path):
    csv_path = os.path.join(base_dir, '..', 'interview_forge_v3_complete.csv')

npy_path = os.path.join(base_dir, 'e5_npu_full_embeddings.npy')
if not os.path.exists(npy_path):
    npy_path = os.path.join(base_dir, 'e5_full_embeddings.npy')
    if not os.path.exists(npy_path):
        npy_path = os.path.join(base_dir, '..', 'e5_npu_full_embeddings.npy')
        if not os.path.exists(npy_path):
            npy_path = os.path.join(base_dir, '..', 'e5_full_embeddings.npy')

df = pd.read_csv(csv_path).dropna(subset=['question']).reset_index(drop=True)
full_embeddings = np.load(npy_path)

final_model = SentenceTransformer("intfloat/e5-small-v2")

model_id = "Qwen/Qwen2.5-1.5B-Instruct"  
print(f"Loading {model_id} into memory...")
generator = pipeline("text-generation", model=model_id, torch_dtype=torch.bfloat16, device="cpu")
print("Models loaded successfully!")

roles = sorted(df['role'].unique().tolist())
sectors = sorted(df['sector'].unique().tolist())
interviewers = ["Strict Technical Lead", "Friendly HR Manager", "Aggressive CISO", "Curious Senior Developer", "Business-Focused Product Manager"]
raw_levels = sorted(df['question_level'].unique().tolist())
levels = [lvl.split(': ')[-1] if ': ' in lvl else lvl for lvl in raw_levels]


def get_interview_question(user_role, user_sector, user_interviewer, user_level):
    query_text = f"An interview question for a {user_role} in the {user_sector} sector focusing on {user_level} concepts, asked by a {user_interviewer}."
    query_embedding = final_model.encode([f"query: {query_text}"], normalize_embeddings=True)
    similarities = cosine_similarity(query_embedding, full_embeddings)[0]
    best_match_idx = similarities.argsort()[::-1][0]
    return df.iloc[best_match_idx]['question']


def get_more_like_this(user_role, user_sector, current_question):
    if not current_question:
        return "Please generate a question first."
    
    filtered_df = df[(df['role'] == user_role) & (df['sector'] == user_sector)]
    
    if filtered_df.empty:
        filtered_df = df 
        
    pool = filtered_df[filtered_df['question'] != current_question]
    
    if pool.empty:
        pool = filtered_df
        
    random_match = pool.sample(n=1).iloc[0]['question']
    return random_match


def get_interview_question_and_clear(*args):
    question = get_interview_question(*args)
    return question, "", "", ""

def get_more_like_this_and_clear(*args):
    question = get_more_like_this(*args)
    return question, "", "", ""

def create_circular_progress(grade_text):
    match = re.search(r'Grade:\s*(\d+)', grade_text)
    if match:
        score = int(match.group(1))
    else:
        score = 0
        
    percentage = (score / 10) * 100
    dasharray = f"{percentage} {100 - percentage}"
    
    if score >= 9:
        color = "#22c55e"  # Bright green
    elif score >= 8:
        color = "#16a34a"  # Darkish green
    elif score >= 6:
        color = "#fb923c"  # Light orange
    elif score >= 4:
        color = "#ea580c"  # Orange
    elif score >= 2:
        color = "#ef4444"  # Red
    else:
        color = "#b91c1c"  # Dark red
        
    svg_html = f"""
    <div style="display: flex; justify-content: center; align-items: center; padding: 20px; flex-direction: column;">
        <h3 style="margin-bottom: 15px; color: #d4af37; font-family: 'Outfit', sans-serif; font-size: 1.5em; font-weight: 600;">AI Grade</h3>
        <div style="position: relative; width: 150px; height: 150px;">
            <svg viewBox="0 0 36 36" style="width: 100%; height: 100%;">
                <path
                    d="M18 2.0845
                      a 15.9155 15.9155 0 0 1 0 31.831
                      a 15.9155 15.9155 0 0 1 0 -31.831"
                    fill="none"
                    stroke="#0b172a"
                    stroke-width="3"
                />
                <path
                    d="M18 2.0845
                      a 15.9155 15.9155 0 0 1 0 31.831
                      a 15.9155 15.9155 0 0 1 0 -31.831"
                    fill="none"
                    stroke="{color}"
                    stroke-width="3"
                    stroke-dasharray="{dasharray}"
                    style="transition: stroke-dasharray 1s ease-out;"
                />
            </svg>
            <div style="position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%); font-size: 28px; font-weight: bold; color: #f8fafc; font-family: 'Outfit', sans-serif;">
                {score}/10
            </div>
        </div>
    </div>
    """
    return svg_html

def evaluate_and_format(question_text, candidate_answer, user_role, user_sector, user_interviewer, user_level):
    if not candidate_answer.strip():
        return "", "Please type an answer before submitting."
        
    system_prompt = f"""You are a {user_interviewer} evaluating a {user_level} {user_role} candidate in the {user_sector} sector.

CRITICAL RULES:
1. Speak DIRECTLY to the candidate using "you" and "your". Never use the word "candidate".
2. Separate the prompt context. Do NOT penalize or critique the user for constraints that were mentioned in the [INTERVIEW QUESTION].
3. You MUST generate an Example Answer at the very end. Keep it extremely short.
4. If the candidate's answer is very short, vague, or a variation of "I don't know", you MUST give a Grade of 1/10 and explicitly state they failed to provide an answer in the Cons.
5. A concise answer is NOT a bad answer. Judge by correctness and relevance, NOT by length.

GRADING SCALE (follow this strictly):
- 9-10: The answer is correct, demonstrates clear understanding, and covers the key points. It does NOT need to be perfect or exhaustive. A real interviewer would be impressed.
- 7-8: The answer is decent but has noticeable gaps in reasoning or misses important concepts.
- 4-6: The answer is partially correct but shows weak understanding or is too surface-level.
- 1-3: The answer is mostly wrong, irrelevant, or the candidate did not attempt it.

IMPORTANT: Only list a Con if it is a genuine mistake or a significant missing concept. Do NOT list "nice-to-have" extras or alternative approaches as Cons. If the answer is strong, give it a 9 or 10.

GRADING EXAMPLES (use these to calibrate your scoring):

Example Question: "How would you secure a REST API?"

Answer: "I'd use HTTPS for encryption in transit, JWT tokens with short expiry for auth, validate and sanitize all inputs, and add rate limiting to prevent abuse." -> Grade: 9/10
Why: Covers the key pillars of API security with specific, correct techniques.

Answer: "I'd start with HTTPS and token-based authentication. I'd also add input validation to prevent injection attacks, though I'm less sure about the best rate limiting approach." -> Grade: 7/10
Why: Solid understanding of core concepts, minor gap is acknowledged honestly.

Answer: "I'd add authentication and maybe some encryption. Also make sure only authorized users can access it." -> Grade: 5/10
Why: Right direction but too vague — no specific techniques or tools mentioned.

Answer: "Probably use passwords and a firewall. Maybe SSL." -> Grade: 3/10
Why: Shows very basic awareness but lacks real understanding of API security.

Answer: "I don't really know, I'd Google it." -> Grade: 1/10
Why: No attempt to answer.

You MUST output exactly this format and nothing else:
Grade: [1-10]/10
Pros:
- [Pro 1]
- [Pro 2]
Cons:
- [Con 1]
- [Con 2]
Example Answer:
[Provide a strict maximum 2-sentence example of a perfect answer.]"""

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"[INTERVIEW QUESTION]\n{question_text}\n\n[CANDIDATE'S ANSWER]\n{candidate_answer}"}
    ]
    
    outputs = generator(messages, max_new_tokens=400, temperature=0.15, do_sample=True)
    raw_feedback = outputs[0]['generated_text'][-1]['content']
    
    # Safety net: prevent unreasonably low grades for real answers
    match = re.search(r'Grade:\s*(\d+)', raw_feedback)
    if match:
        model_score = int(match.group(1))
        word_count = len(candidate_answer.split())
        
        # If the answer has real substance, enforce a minimum floor
        if word_count >= 40:
            min_grade = 4
        elif word_count >= 20:
            min_grade = 3
        elif word_count >= 8:
            min_grade = 2
        else:
            min_grade = 1
        
        if model_score < min_grade:
            adjusted_score = min_grade
            raw_feedback = re.sub(r'Grade:\s*\d+', f'Grade: {adjusted_score}', raw_feedback)
    
    score_html = create_circular_progress(raw_feedback)
    clean_feedback = re.sub(r'Grade:.*?\n', '', raw_feedback).strip()
    
    return score_html, clean_feedback

# =====================================================================
# UI DESIGN & THEME INJECTION
# =====================================================================

custom_css = """
body, .gradio-container {
    background-color: #040f23 !important;
}
.centered-dropdowns {
    max-width: 900px !important;
    margin: 0 auto !important;
}
.center-btn {
    max-width: 300px !important;
    margin: 20px auto !important;
    display: block !important;
    font-size: 1.2em !important;
    font-weight: bold !important;
}



/* Smaller font inside the Question and Answer textboxes */
textarea {
    font-size: 0.88em !important;
    line-height: 1.6 !important;
}
.side-by-side {
    align-items: stretch !important;
    gap: 20px !important;
}

/* Fix Dropdown Menu to be Dark and Readable */
.options, .options-wrap, [role="listbox"], ul.options {
    background-color: #1e293b !important;
    color: #f8fafc !important;
    border: 1px solid #475569 !important;
}
li.item {
    background-color: #1e293b !important;
    color: #f8fafc !important;
}
li.item:hover, li.item.selected {
    background-color: #334155 !important;
    color: #d4af37 !important;
}

/* Make Dropdown and Textbox Inputs brighter for contrast */
textarea, input {
    background-color: #1e293b !important;
    color: #f8fafc !important;
    border: 1px solid #475569 !important;
    border-radius: 6px !important;
}

/* Aggressively Force ALL Headers/Labels to be Golden and Large */
label span, .block-title, label, label.svelte-1b6s6s span.svelte-1b6s6s, .label-text, .form-label {
    font-size: 1.45em !important;
    color: #d4af37 !important;
    font-family: 'Outfit', sans-serif !important;
    font-weight: 700 !important;
    margin-bottom: 10px !important;
    display: block !important;
    letter-spacing: 0.02em !important;
}
label * {
    color: #d4af37 !important;
    font-size: 1.45em !important;
    font-weight: 700 !important;
}

/* Override the CSS VARIABLES that Gradio's Svelte styles read from */
#dd-role, #dd-sector, #dd-interviewer, #dd-level {
    --block-title-text-color: #d4af37 !important;
    --block-title-text-size: 1.3em !important;
    --block-title-text-weight: 700 !important;
    --block-label-text-color: #d4af37 !important;
}
#dd-role span, #dd-sector span, #dd-interviewer span, #dd-level span {
    color: #d4af37 !important;
}
"""

theme = gr.themes.Default(
    font=(gr.themes.GoogleFont("Outfit"), "sans-serif"),
).set(
    body_background_fill="#040f23",
    body_background_fill_dark="#040f23",
    body_text_color="#f8fafc",
    body_text_color_dark="#f8fafc",
    
    background_fill_primary="#040f23",
    background_fill_primary_dark="#040f23",
    background_fill_secondary="#1e293b",
    background_fill_secondary_dark="#1e293b",
    
    block_background_fill="#040f23", 
    block_background_fill_dark="#040f23",
    block_border_width="0px",
    
    # Textboxes / Dropdowns Base (Backed up by CSS)
    input_background_fill="#1e293b", 
    input_background_fill_dark="#1e293b",
    input_border_color="#475569", 
    input_border_color_dark="#475569",
    input_border_width="1px",
    
    # Labels Base
    block_label_text_color="#d4af37", 
    block_label_text_color_dark="#d4af37",
    
    # Buttons
    button_primary_background_fill="#d4af37", 
    button_primary_background_fill_dark="#d4af37",
    button_primary_text_color="#000000",
    button_primary_text_color_dark="#000000",
    
    button_secondary_background_fill="#1e293b",
    button_secondary_background_fill_dark="#1e293b",
    button_secondary_text_color="#f8fafc",
    button_secondary_text_color_dark="#f8fafc"
)

# Load Logo safely via Base64 so it NEVER fails on Hugging Face Spaces
logo_path = os.path.join(base_dir, "logo.png")
if not os.path.exists(logo_path):
    logo_path = os.path.join(base_dir, "..", "logo.png")

if os.path.exists(logo_path):
    with open(logo_path, "rb") as img_file:
        b64_string = base64.b64encode(img_file.read()).decode('utf-8')
    logo_html = f'<img src="data:image/png;base64,{b64_string}" style="max-height: 250px; margin: 0 auto; display: block;" alt="Interview Forge"/>'
else:
    logo_html = '<h1 style="color: #d4af37; text-align: center; font-size: 3em;">Interview Forge</h1><p style="text-align: center; color: #94a3b8;">Success Through Preparation</p>'


with gr.Blocks(theme=theme, css=custom_css) as app:
    
    # Render the Base64 Logo perfectly
    gr.HTML(f"""<div style="text-align: center; margin-bottom: 30px; margin-top: 10px;">{logo_html}</div>""")

    gr.HTML("<h2 style='text-align: center; color: #d4af37; font-family: Outfit, sans-serif; font-size: 1.8em; font-weight: 700; margin-bottom: 16px; letter-spacing: 0.03em;'>Quick Start</h2>")
    with gr.Row(elem_classes="centered-dropdowns"):
        starter_1 = gr.Button("Data Scientist (FinTech)", variant="secondary")
        starter_2 = gr.Button("Backend Developer (Cybersecurity)", variant="secondary")
        starter_3 = gr.Button("UX/UI Designer (SaaS)", variant="secondary")
    gr.Markdown("<br>")
    
    with gr.Column(elem_classes="centered-dropdowns"):
        with gr.Row():
            role_dropdown = gr.Dropdown(choices=roles, label="Role", value=roles[0] if roles else None, elem_id="dd-role")
            sector_dropdown = gr.Dropdown(choices=sectors, label="Sector", value=sectors[0] if sectors else None, elem_id="dd-sector")
        with gr.Row():
            interviewer_dropdown = gr.Dropdown(choices=interviewers, label="Interviewer Persona", value=interviewers[0], elem_id="dd-interviewer")
            level_dropdown = gr.Dropdown(choices=levels, label="Difficulty Level", value=levels[1], elem_id="dd-level")
            
    generate_btn = gr.Button("Generate Custom Question", variant="primary", elem_classes="center-btn")
    
    with gr.Row(elem_classes="side-by-side"):
        with gr.Column():
            question_display = gr.Textbox(label="The Question", interactive=False, lines=8, text_align="left")
            more_btn = gr.Button("More Like This", variant="secondary")
        with gr.Column():
            user_answer = gr.Textbox(label="Your Answer", lines=8, placeholder="Type your answer here...")
            submit_btn = gr.Button("Submit Answer for AI Grading", variant="primary")
            
    with gr.Row():
        with gr.Column(scale=1, min_width=200):
            score_circle = gr.HTML()
        with gr.Column(scale=3):
            feedback_display = gr.Markdown()

    generate_btn.click(
        fn=get_interview_question_and_clear, 
        inputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown], 
        outputs=[question_display, user_answer, score_circle, feedback_display]
    )
    
    more_btn.click(
        fn=get_more_like_this_and_clear,
        inputs=[role_dropdown, sector_dropdown, question_display],
        outputs=[question_display, user_answer, score_circle, feedback_display]
    )

    # Quick Starter Events
    starter_1.click(fn=lambda: ("Data Scientist", "FinTech", "Strict Technical Lead", "Practical"), outputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown]).then(
        fn=get_interview_question_and_clear,
        inputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown],
        outputs=[question_display, user_answer, score_circle, feedback_display]
    )
    starter_2.click(fn=lambda: ("Backend Developer", "Cybersecurity", "Curious Senior Developer", "Foundational"), outputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown]).then(
        fn=get_interview_question_and_clear,
        inputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown],
        outputs=[question_display, user_answer, score_circle, feedback_display]
    )
    starter_3.click(fn=lambda: ("UX/UI Designer", "SaaS & Cloud Platforms", "Business-Focused Product Manager", "Edge Case & Conflict"), outputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown]).then(
        fn=get_interview_question_and_clear,
        inputs=[role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown],
        outputs=[question_display, user_answer, score_circle, feedback_display]
    )
    
    submit_btn.click(
        fn=evaluate_and_format, 
        inputs=[question_display, user_answer, role_dropdown, sector_dropdown, interviewer_dropdown, level_dropdown], 
        outputs=[score_circle, feedback_display]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, share=False)