import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import re
import json
from datetime import datetime

# ── Load model ───────────────────────────────────────────────
print("Loading base model...")
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
ADAPTER    = "likithyadavv/codementor-7b"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    dtype=torch.float16,
    trust_remote_code=True,
    low_cpu_mem_usage=True
)

model = PeftModel.from_pretrained(base_model, ADAPTER, is_trainable=False)
model = model.merge_and_unload()

device = "cuda" if torch.cuda.is_available() else "cpu"
model  = model.to(device)
model.eval()
print(f"✅ Model ready on {device}!")

# ── Language detection ───────────────────────────────────────
def detect_language(text):
    text_lower = text.lower()
    langs = {
        "Java":       ["java", "system.out", "public class", "import java", "scanner"],
        "Python":     ["python", "def ", "print(", "elif", "range("],
        "C++":        ["c++", "cout", "#include", "cin", "namespace"],
        "C":          ["printf", "scanf", "#include <stdio"],
        "JavaScript": ["console.log", "const ", "let ", "function "],
    }
    for lang, keywords in langs.items():
        if any(kw in text_lower for kw in keywords):
            return lang
    return "General"

# ── Code detection ───────────────────────────────────────────
def has_code(text):
    patterns = [
        r"public\s+class", r"def\s+\w+\(",
        r"#include",        r"import\s+java",
        r"System\.out\.print", r"while\s*\(",
        r"for\s*\(",        r"console\.log",
        r"```",             r"int\s+\w+\s*=",
        r"for\s+\w+\s+in\s+"
    ]
    return any(re.search(p, text) for p in patterns)

# ── Detect student signal ────────────────────────────────────
def detect_student_signal(text):
    text_lower = text.lower()
    if any(k in text_lower for k in ["hello", "hi", "hey", "good morning", "good evening", "howdy"]):
        return "greeting"
    if any(k in text_lower for k in ["hint", "help", "don't know", "idk", "no idea", "i give up", "tell me", "just give"]):
        return "hint_request"
    if any(k in text_lower for k in ["i think", "maybe", "is it", "could it be", "how about", "like this", "would it"]):
        return "attempting"
    if any(k in text_lower for k in ["what is", "explain", "don't understand", "confused", "what does", "difference"]):
        return "concept_question"
    if has_code(text):
        return "submitted_code"
    return "general"

# ── Extract content from a history message (dict or ChatMessage) ──
def msg_role(m):
    return m["role"] if isinstance(m, dict) else m.role

def msg_content(m):
    c = m["content"] if isinstance(m, dict) else m.content
    # If content is itself a dict (ChatMessage serialization artifact), unwrap it
    if isinstance(c, dict):
        return c.get("text", c.get("content", str(c)))
    return str(c)

# ── Build adaptive prompt ────────────────────────────────────
def build_prompt(user_message, history, language):
    signal = detect_student_signal(user_message)

    history_text = ""
    if history:
        pairs = []
        i = 0
        while i < len(history) - 1:
            a, b = history[i], history[i + 1]
            if msg_role(a) == "user" and msg_role(b) == "assistant":
                pairs.append((msg_content(a), msg_content(b)))
                i += 2
            else:
                i += 1
        for human, assistant in pairs[-3:]:
            history_text += f"Student: {human}\nCodeMentor: {assistant}\n\n"

    context = f"Previous conversation:\n{history_text}" if history_text else ""

    styles = {
        "greeting":         "The student just said hello.\n- Greet them briefly and warmly (1 sentence max)\n- Ask what programming topic or problem they want to work on today\n- Do NOT give any code or explanations yet",
        "hint_request":     "The student is asking for a hint.\n- Give ONE small hint only — not the answer\n- Point them in the right direction with a question\n- Do NOT write complete code",
        "submitted_code":   "The student has written code.\n- First check if the logic is correct or not\n- If CORRECT: validate genuinely then push forward with a question\n- If WRONG: identify the specific issue, give a targeted hint\n- If copy-pasted (too perfect): ask them to walk you through it\n- Never rewrite their code for them",
        "attempting":       "The student is guessing or attempting an explanation.\n- If CORRECT: confirm clearly and move to the next concept\n- If WRONG: gently correct their thinking with a question\n- Always build on what they said",
        "concept_question": "The student is asking about a concept.\n- Start with a simple real-world analogy\n- Give a tiny example with _____ blanks for them to fill\n- End with a question to test understanding\n- No textbook definitions",
        "general":          "General question.\n- Understand what they are trying to build or learn\n- Give the first small step only\n- Never write the full solution",
    }
    style = styles.get(signal, styles["general"])

    return f"""### System:
You are CodeMentor — a sharp, patient programming tutor for Python, Java, C, and C++.
Your only job is to move the student forward one step at a time.
You NEVER write complete solutions.
You always respond to exactly what the student gave you.
Detected language: {language}
{style}

{context}### Instruction:
Respond to the student's message below.

### Input:
{user_message}

### Response:
CodeMentor:"""

# ── Extract clean response ───────────────────────────────────
def extract_response(full_text):
    if "### Response:" in full_text:
        after = full_text.split("### Response:")[-1]
    else:
        after = full_text

    if after.lstrip().startswith("CodeMentor:"):
        after = after.lstrip()[len("CodeMentor:"):].strip()
    else:
        after = after.strip()

    for stop in ["### Input:", "### Instruction:", "Student:", "\n###"]:
        if stop in after:
            after = after.split(stop)[0]

    return after.strip()

# ── Generate response ────────────────────────────────────────
def generate(prompt):
    device = next(model.parameters()).device
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.2,
            pad_token_id=tokenizer.eos_token_id,
        )
    full = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return extract_response(full)

# ── Log conversations ────────────────────────────────────────
def log_conversation(user_msg, bot_msg, language):
    log = {"timestamp": datetime.now().isoformat(), "language": language, "user": user_msg, "assistant": bot_msg}
    with open("conversations.jsonl", "a") as f:
        f.write(json.dumps(log) + "\n")

# ── Main chat function ───────────────────────────────────────
def chat(user_message, history):
    language = detect_language(user_message)
    if language == "General" and history:
        for m in history:
            if msg_role(m) == "user":
                detected = detect_language(msg_content(m))
                if detected != "General":
                    language = detected
                    break

    response = generate(build_prompt(user_message, history, language))
    log_conversation(user_message, response, language)
    return response, language

# ── Gradio UI ────────────────────────────────────────────────
with gr.Blocks(title="CodeMentor AI") as demo:
    gr.Markdown("# 🎓 CodeMentor — AI Programming Tutor")
    gr.Markdown("*Ask any programming question. I'll guide you — not give you the answer!*")

    chatbot = gr.Chatbot(height=450)

    lang_display = gr.Textbox(label="Detected Language", interactive=False, value="🌐 General")

    with gr.Row():
        msg  = gr.Textbox(placeholder="e.g. How do I write a loop in Java?", label="Your Question", scale=4)
        send = gr.Button("Ask Mentor 🎓", variant="primary", scale=1)

    clear = gr.Button("🗑 Clear Chat")

    def respond(message, chat_history):
        if not message.strip():
            return "", chat_history, lang_display.value
        response, language = chat(message, chat_history)
        # Plain dicts — the correct format for this Gradio version
        chat_history.append({"role": "user",      "content": message})
        chat_history.append({"role": "assistant", "content": response})
        return "", chat_history, f"🌐 {language}"

    send.click(respond,  [msg, chatbot], [msg, chatbot, lang_display])
    msg.submit(respond,  [msg, chatbot], [msg, chatbot, lang_display])
    clear.click(lambda: ([], "🌐 General"), outputs=[chatbot, lang_display])

demo.launch(theme=gr.themes.Soft())