import gradio as gr
from transformers import pipeline

# =====================================
# Load Model
# =====================================

generator = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-1.5B-Instruct"
)

# =====================================
# Master Prompt
# =====================================

MASTER_PROMPT = """
You are an expert technical interviewer.

Generate interview questions for the candidate below.

Role: {role}
Experience: {experience}
Skills: {skills}
Company: {company}

IMPORTANT:

Return your response EXACTLY in the following format.

<<<HR>>>
1.
2.
3.
4.
5.

<<<TECHNICAL>>>
1.
2.
3.
4.
5.

<<<CODING>>>
1.
2.
3.
4.
5.

<<<BEHAVIORAL>>>
1.
2.
3.
4.
5.

<<<OTHER>>>
1.
2.
3.
4.
5.

Rules

HR
- Questions about communication, teamwork, strengths, weaknesses, leadership and culture fit.

TECHNICAL
- Conceptual questions only.
- No coding problems.

CODING
- Programming problems only.
- No solutions.
- Easy to Medium difficulty.

BEHAVIORAL
- Scenario based questions.
- Leadership.
- Conflict resolution.
- Time management.
- Decision making.

OTHER
- Career goals.
- Industry awareness.
- Ethics.
- Innovation.
- Learning habits.

Generate exactly FIVE questions for each section.

Do not write anything outside the markers.

Do not use markdown.

Do not add introductions.

Do not add conclusions.
"""

# =====================================
# Split Output
# =====================================

def split_response(text):

    sections = {
        "HR": "",
        "TECHNICAL": "",
        "CODING": "",
        "BEHAVIORAL": "",
        "OTHER": ""
    }

    current = None

    for line in text.splitlines():

        line = line.strip()

        if line == "<<<HR>>>":
            current = "HR"
            continue

        elif line == "<<<TECHNICAL>>>":
            current = "TECHNICAL"
            continue

        elif line == "<<<CODING>>>":
            current = "CODING"
            continue

        elif line == "<<<BEHAVIORAL>>>":
            current = "BEHAVIORAL"
            continue

        elif line == "<<<OTHER>>>":
            current = "OTHER"
            continue

        if current:
            sections[current] += line + "\n"

    return (
        sections["HR"].strip(),
        sections["TECHNICAL"].strip(),
        sections["CODING"].strip(),
        sections["BEHAVIORAL"].strip(),
        sections["OTHER"].strip()
    )

# =====================================
# Generate Questions
# =====================================

def generate_questions(role, experience, skills, company):

    if not company.strip():
        company = "Not Specified"

    prompt = MASTER_PROMPT.format(
        role=role,
        experience=experience,
        skills=skills,
        company=company
    )

    messages = [
        {
            "role": "system",
            "content": "You are an expert interviewer."
        },
        {
            "role": "user",
            "content": prompt
        }
    ]

    chat_prompt = generator.tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    output = generator(
        chat_prompt,
        max_new_tokens=700,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        return_full_text=False
    )

    response = output[0]["generated_text"]

    return split_response(response)

# =====================================
# Simple Gradio UI
# =====================================

with gr.Blocks(title="Interview Preparation Assistant") as demo:

    gr.Markdown("# 🎯 Interview Preparation Assistant")
    gr.Markdown(
        "Generate **HR, Technical, Coding, Behavioral, and Other** interview questions."
    )

    role = gr.Textbox(
        label="Job Role",
        placeholder="Software Engineer"
    )

    experience = gr.Dropdown(
        choices=[
            "Fresher",
            "0-2 Years",
            "2-5 Years",
            "5+ Years"
        ],
        value="Fresher",
        label="Experience"
    )

    skills = gr.Textbox(
        label="Skills",
        placeholder="Python, SQL, Machine Learning, DSA"
    )

    company = gr.Textbox(
        label="Target Company (Optional)",
        placeholder="Google, Microsoft, Amazon..."
    )

    generate_btn = gr.Button(
        "Generate Interview Questions",
        variant="primary"
    )

    gr.Markdown("---")

    hr_box = gr.Textbox(
        label="HR Questions",
        lines=12
    )

    technical_box = gr.Textbox(
        label="Technical Questions",
        lines=12
    )

    coding_box = gr.Textbox(
        label="Coding Questions",
        lines=12
    )

    behavioral_box = gr.Textbox(
        label="Behavioral Questions",
        lines=12
    )

    other_box = gr.Textbox(
        label="Other Questions",
        lines=12
    )

    generate_btn.click(
        fn=generate_questions,
        inputs=[
            role,
            experience,
            skills,
            company
        ],
        outputs=[
            hr_box,
            technical_box,
            coding_box,
            behavioral_box,
            other_box
        ]
    )

if __name__ == "__main__":
    demo.launch()