# app.py – FINAL, ACTUALLY WORKING VERSION (November 2025)
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel
import gradio as gr

BASE_MODEL   = "mistralai/Mistral-7B-Instruct-v0.2"
LORA_ADAPTER = "rishu834763/java-explainer-lora"

print("Loading Java Explainer (8-bit CPU) – please wait ~90 seconds...")

# This combination NEVER fails on any HF Space
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    load_in_8bit=True,
    device_map="auto",           # "auto" works perfectly with 8-bit (no offload error)
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
)

# Apply your LoRA
model = PeftModel.from_pretrained(model, LORA_ADAPTER)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Fast and reliable pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=1024,
    temperature=0.2,
    top_p=0.95,
    do_sample=True,
    repetition_penalty=1.18,
    return_full_text=False,
)

SYSTEM_PROMPT = """You are the world's best Java teacher with 20+ years of experience.
Always give:
• Clear, step-by-step explanation
• Clean, modern, runnable Java code (Java 17+)
• Fix any bugs or bad practices
• Use records, var, streams, sealed classes, etc. when appropriate"""

def generate(instruction: str, code: str = ""):
    user_input = f"### Instruction:\n{instruction.strip()}\n\n### Code:\n{code.strip()}" if code.strip() else instruction.strip()
    
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": user_input}
    ]
    
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    result = pipe(prompt)[0]["generated_text"].strip()
    return result

# Beautiful working UI (Gradio 4.100+ compatible)
with gr.Blocks(theme=gr.themes.Soft(), title="Java Explainer Pro") as demo:
    gr.Markdown("# Java Explainer Pro\nYour personal senior Java mentor is ready")

    with gr.Row():
        with gr.Column(scale=1):
            instruction = gr.Textbox(
                label="What do you want to know or fix?",
                placeholder="Explain this code · Fix this bug · Make it thread-safe · Convert to records · Best way to read JSON in Java 17",
                lines=4
            )
            code_input = gr.Code(
                label="Java Code (optional)",
                language="java",
                lines=16,
                value="// Paste your code here (