import gradio as gr
import torch
import re
import random
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
from datasets import load_dataset
model     = None
tokenizer = None
eval_data = None
QUESTION_COL = None
ANSWER_COL   = None

BASE_MODEL_ID    = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER_PATH     = "./qlora_interview_adapter"   # your saved adapter folder
#  LOAD MODEL  
def load_model():
    global model, tokenizer, eval_data, QUESTION_COL, ANSWER_COL

    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        ADAPTER_PATH if os.path.exists(ADAPTER_PATH) else BASE_MODEL_ID
    )
    tokenizer.pad_token = tokenizer.eos_token

    # 4-bit QLoRA config 
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16
    )

    print("Loading base model in 4-bit...")
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto"
    )

    # Load your fine-tuned LoRA adapter if it exists
    if os.path.exists(ADAPTER_PATH):
        print(f"Loading QLoRA adapter from {ADAPTER_PATH}...")
        model = PeftModel.from_pretrained(base, ADAPTER_PATH)
        print("QLoRA adapter loaded!")
    else:
        print("No adapter found  using base TinyLlama ")
        model = base

    model.eval()

    # Load dataset for question pool
    print("Loading interview dataset...")
    dataset = load_dataset("Aiman1234/Interview-questions", split="train")
    dataset = dataset.select(range(min(500, len(dataset))))

    cols = dataset.column_names
    for c in cols:
        lc = c.strip().lower()
        if "question" in lc: QUESTION_COL = c
        if "answer"   in lc: ANSWER_COL   = c

    split     = dataset.train_test_split(test_size=0.2, seed=42)
    eval_data = split["test"]
    print(f"Ready! {len(eval_data)} eval questions loaded.")

#  SAME FUNCTIONS FROM YOUR NOTEBOOK

def generate_model_answer(question, max_new_tokens=300):
    """Step 10 from your notebook — generate reference answer from fine-tuned model."""
    prompt = f"### Instruction:\n{question}\n\n### Response:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.3,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )

    new_tokens = output[0][inputs['input_ids'].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def compute_ngram_similarity(text1, text2, n=2):
    """BLEU-inspired n-gram overlap."""
    def get_ngrams(text, n):
        words = re.findall(r'\w+', text.lower())
        return set(zip(*[words[i:] for i in range(n)]))

    ngrams1 = get_ngrams(text1, n)
    ngrams2 = get_ngrams(text2, n)

    if not ngrams1 or not ngrams2:
        return 0.0

    overlap   = ngrams1 & ngrams2
    precision = len(overlap) / len(ngrams1)
    recall    = len(overlap) / len(ngrams2)

    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


def llm_evaluate(question, model_answer, user_answer, max_new_tokens=250):
    """ LLM-based evaluation using the fine-tuned model."""
    eval_prompt = f"""### Instruction:
You are an AI interview evaluator. Compare the user's answer with the expected answer.

Question: {question}

Expected Answer: {model_answer[:400]}

User Answer: {user_answer}
"""
    inputs = tokenizer(eval_prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.2,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )

    new_tokens = output[0][inputs['input_ids'].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def parse_evaluation(llm_output, ngram_score):
    """ parse score and feedback."""
    score_match    = re.search(r'Score:\s*(\d+(?:\.\d+)?)/10', llm_output, re.IGNORECASE)
    feedback_match = re.search(r'Feedback:\s*(.+)',            llm_output, re.IGNORECASE | re.DOTALL)

    score    = float(score_match.group(1)) if score_match else round(ngram_score * 10, 1)
    feedback = feedback_match.group(1).split('\n')[0].strip() if feedback_match else (
        "Good answer!" if ngram_score >= 0.7 else
        "Partial answer. Consider expanding on key points." if ngram_score >= 0.4 else
        "Needs improvement. Review key concepts."
    )

    return max(0.0, min(10.0, score)), feedback


def evaluate_answer_fast(user_answer, dataset_answer):
    """
    Fast evaluation using only n-gram similarity against the dataset answer.
    No model inference — instant results.
    """
    ngram_sim_1 = compute_ngram_similarity(user_answer, dataset_answer, n=1)
    ngram_sim_2 = compute_ngram_similarity(user_answer, dataset_answer, n=2)
    ngram_score = (ngram_sim_1 + ngram_sim_2) / 2

    score = round(ngram_score * 10, 1)

    if ngram_score >= 0.7:
        feedback = "Great answer! Strong overlap with the reference answer."
    elif ngram_score >= 0.4:
        feedback = "Partial answer. Consider expanding on key points from the reference."
    else:
        feedback = "Needs improvement. Review the reference answer and key concepts."

    return score, feedback, round(ngram_score * 100, 1)
#  SESSION STATE HELPERS
def get_random_questions(n=5):
    indices = random.sample(range(len(eval_data)), min(n, len(eval_data)))
    questions       = [eval_data[i][QUESTION_COL] for i in indices]
    dataset_answers = [eval_data[i][ANSWER_COL]   for i in indices]
    return questions, dataset_answers


def start_session(num_q):
    """Pick questions, return first one."""
    questions, dataset_answers = get_random_questions(int(num_q))
    state = {
        "questions":       questions,
        "dataset_answers": dataset_answers,
        "current":         0,
        "results":         [],
        "total":           int(num_q)
    }
    q = questions[0]
    progress = f"Question 1 / {num_q}"
    return state, q, progress, "", "", "", gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)


def submit_answer(user_answer, state):
    """Evaluate current answer, show result."""
    if not user_answer.strip():
        return state, "⚠️ Please type your answer first.", "", "", gr.update(visible=False)

    idx            = state["current"]
    q              = state["questions"][idx]
    dataset_answer = state["dataset_answers"][idx]

    score, feedback, ngram_pct = evaluate_answer_fast(user_answer, dataset_answer)

    state["results"].append({
        "question":    q,
        "user_answer": user_answer,
        "score":       score,
        "feedback":    feedback,
        "ngram_pct":   ngram_pct
    })

    score_display = (
        f"## Score: {score:.1f} / 10\n\n"
        f"**Feedback:** {feedback}\n\n"
        f"**N-gram Similarity (BLEU-style):** {ngram_pct}%\n\n"
        f"---\n\n"
        f"###  Reference Answer\n"
        f"{dataset_answer}"
    )
    is_last    = state["current"] >= state["total"] - 1
    next_label = "See Final Results →" if is_last else "Next Question →"

    return state, "", score_display, next_label, gr.update(visible=True)


def next_question(state):
    """Move to next question or show summary."""
    state["current"] += 1

    if state["current"] >= state["total"]:
        # Build summary
        results   = state["results"]
        avg_score = sum(r["score"] for r in results) / len(results)

        if avg_score >= 8.5: verdict = "🏆 Outstanding — You are very well prepared!"
        elif avg_score >= 7: verdict = "✅ Strong performance — A few areas to polish."
        elif avg_score >= 5.5: verdict = "📈 Good effort — Review feedback to improve."
        elif avg_score >= 4: verdict = "📚 Keep studying — Focus on key concepts."
        else: verdict = "💪 Don't give up — Practice makes perfect!"

        summary = f"## Session Complete!\n\n**Average Score: {avg_score:.1f} / 10**\n\n{verdict}\n\n---\n\n"
        for i, r in enumerate(results, 1):
            summary += f"**Q{i}:** {r['question']}\n"
            summary += f"- Score: {r['score']:.1f}/10 | N-gram: {r['ngram_pct']}%\n"
            summary += f"- Feedback: {r['feedback']}\n\n"

        return (state, "Session complete!", f"0 / {state['total']}",
                gr.update(visible=False), gr.update(visible=True),
                summary, gr.update(visible=False))

    q        = state["questions"][state["current"]]
    progress = f"Question {state['current']+1} / {state['total']}"

    return (state, q, progress,
            gr.update(visible=True), gr.update(visible=False),
            "", gr.update(visible=False))

#  GRADIO UI
with gr.Blocks(title="AI Interview Evaluator — QLoRA") as demo:

    state = gr.State({})

    gr.Markdown("""
    #  AI Interview Evaluator
    """, elem_classes="title-block")

    # ── SETUP PANEL ──
    with gr.Group(visible=True) as setup_panel:
        gr.Markdown("### Start a new session")
        num_q = gr.Dropdown(
            choices=["3","5","8","10"],
            value="5",
            label="Number of Questions"
        )
        start_btn = gr.Button(" Start Interview Session", variant="primary", size="lg")

    # ── INTERVIEW PANEL ──
    with gr.Group(visible=False) as interview_panel:
        progress_label = gr.Markdown("Question 1 / 5")

        question_box = gr.Textbox(
            label="📋 Question",
            interactive=False,
            lines=3,
            elem_classes="question-box"
        )

        answer_box = gr.Textbox(
            label="✍️ Your Answer",
            placeholder="Type your answer here... take your time.",
            lines=6
        )

        with gr.Row():
            submit_btn = gr.Button("Evaluate My Answer ✅", variant="primary")
            next_btn   = gr.Button("Next Question →", variant="secondary", visible=False)

        result_box = gr.Markdown("", label="Evaluation Result")

    # ── SUMMARY PANEL ──
    with gr.Group(visible=False) as summary_panel:
        summary_box  = gr.Markdown("")
        restart_btn  = gr.Button("🔄 Start New Session", variant="primary")

    # ── WIRING ──
    start_btn.click(
        fn=start_session,
        inputs=[num_q],
        outputs=[state, question_box, progress_label, answer_box, result_box,
                 result_box, setup_panel, interview_panel, summary_panel]
    )

    submit_btn.click(
        fn=submit_answer,
        inputs=[answer_box, state],
        outputs=[state, answer_box, result_box, next_btn, next_btn]
    )

    next_btn.click(
        fn=next_question,
        inputs=[state],
        outputs=[state, question_box, progress_label,
                 interview_panel, summary_panel,
                 summary_box, next_btn]
    )

    restart_btn.click(
        fn=lambda: (gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)),
        outputs=[setup_panel, interview_panel, summary_panel]
    )
#  STARTUP
if __name__ == "__main__":
    load_model()
    demo.launch(
    css="""
        footer { display: none !important; }
        .title-block { text-align: center; }
    """
)