import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ---------------------------------
# CONFIG
# ---------------------------------
# BASE_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
LORA_PATH = "vastu_lora_adapter_975"
DEVICE = "cpu"

HF_TOKEN = os.getenv("HF_TOKEN")  # 🔥 REQUIRED

SYSTEM_PROMPT = """You are a strict and authoritative Vastu Shastra expert.
You clearly classify every placement as IDEAL, ACCEPTABLE, or INADVISABLE.
You always give practical remedies if something is wrong.
Your tone is confident, traditional, and decisive.
Avoid unnecessary philosophy. Be precise and actionable.
"""

# ---------------------------------
# LOAD MODEL (CPU SAFE)
# ---------------------------------
@torch.inference_mode()
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL,
        token=HF_TOKEN,
        trust_remote_code=True,
    )
    tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        token=HF_TOKEN,                # 🔥 REQUIRED
        torch_dtype=torch.float32,
        device_map="cpu",              # 🔥 FIXED
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )

    model = PeftModel.from_pretrained(
        base_model,
        LORA_PATH,
    )

    model.eval()
    return model, tokenizer


model, tokenizer = load_model()

# ---------------------------------
# GENERATION (CPU OPTIMIZED)
# ---------------------------------
def generate_response(user_prompt):
    prompt = f"""### System:
{SYSTEM_PROMPT}
### User:
{user_prompt}
### Response:
"""

    inputs = tokenizer(prompt, return_tensors="pt")

    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        temperature=0.3,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
    )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return decoded.split("### Response:")[-1].strip()


# ---------------------------------
# GRADIO UI
# ---------------------------------
with gr.Blocks(title="🧭 Vastu AI Advisor (CPU)") as demo:
    gr.Markdown("# 🧭 Vastu AI Advisor")
    gr.Markdown(
        "**CPU-based demo.** Responses may take ~30–45 seconds.\n\n"
        "Ask questions about room placement, directions, and remedies."
    )

    chatbot = gr.Chatbot(height=420)
    msg = gr.Textbox(
        placeholder="Ask a Vastu question (e.g. Is a toilet in NE acceptable?)",
        lines=2,
    )

    def chat(user_msg, history):
        response = generate_response(user_msg)
        history.append((user_msg, response))
        return history, ""

    msg.submit(chat, [msg, chatbot], [chatbot, msg])


demo.launch(
    server_name="0.0.0.0",   # 🔥 REQUIRED FOR HF SPACES
    server_port=7860,
)