import spaces  # MUST be imported before torch

import time
import json
import logging
import uuid
from datetime import datetime, timezone

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


# ============================================================
# Logging setup (monitoring)
# ============================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger("mediguide")

REQUEST_LOG_FILE = "request_log.jsonl"  # ephemeral storage on free Spaces, resets on restart


def log_request(request_id, message, response, latency_s, flagged, error=None):
    entry = {
        "request_id": request_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "message_len": len(message),
        "response_len": len(response) if response else 0,
        "latency_s": round(latency_s, 3),
        "flagged": flagged,
        "error": error,
    }
    logger.info(json.dumps(entry))
    try:
        with open(REQUEST_LOG_FILE, "a") as f:
            f.write(json.dumps(entry) + "\n")
    except Exception as e:
        logger.warning(f"Could not write to request log file: {e}")


# ============================================================
# Guardrails
# ============================================================

EMERGENCY_KEYWORDS = [
    "suicide", "kill myself", "want to die", "end my life",
    "chest pain", "can't breathe", "cannot breathe", "heart attack",
    "overdose", "severe bleeding", "unconscious", "stroke symptoms",
]

DISCLAIMER = (
    "\n\n⚠️ This is not a substitute for professional medical advice. "
    "If this is urgent, please contact a doctor or emergency services."
)

EMERGENCY_RESPONSE = (
    "This sounds like it could be a medical emergency. Please contact "
    "your local emergency services or go to the nearest emergency room "
    "immediately. I'm not able to provide emergency medical care."
)


def check_emergency(message: str) -> bool:
    msg_lower = message.lower()
    return any(kw in msg_lower for kw in EMERGENCY_KEYWORDS)


def apply_output_guardrails(response: str) -> str:
    if "not a substitute" not in response.lower() and "professional medical" not in response.lower():
        response = response.strip() + DISCLAIMER
    return response


# ============================================================
# Configuration
# ============================================================

BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER_MODEL = "rolmaxx/MediGuide-QLoRA"

SYSTEM_PROMPT = """You are MediGuide, a medical conversational assistant.
Provide clear, helpful, and appropriately cautious responses.
Do not claim to replace a qualified healthcare professional.
If symptoms may indicate an emergency, advise the user to seek
appropriate professional medical care."""


# ============================================================
# Tokenizer & Model – loaded on CPU (no GPU yet)
# ============================================================

logger.info("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

logger.info("Loading base model on CPU...")
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="cpu",
)
logger.info("Base model loaded on CPU.")

logger.info("Loading MediGuide QLoRA adapter on CPU...")
model = PeftModel.from_pretrained(
    model,
    ADAPTER_MODEL,
    torch_device="cpu",
)
model = model.merge_and_unload()
model.eval()
logger.info("MediGuide loaded successfully (on CPU).")


# ============================================================
# Core generation
# ============================================================

@spaces.GPU(duration=60)
def _run_model(messages):
    model.to("cuda")

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

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=4096,
    )
    inputs = {k: v.to("cuda") for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
    response = tokenizer.decode(generated_tokens, skip_special_tokens=True)

    model.to("cpu")
    torch.cuda.empty_cache()

    return response.strip()


def generate_response(message, history):
    """
    history: list of {"role": ..., "content": ...} dicts (Gradio 'messages' format,
    also used as the request body shape for the /chat API below).
    """
    request_id = str(uuid.uuid4())[:8]
    start = time.time()

    if not isinstance(message, str):
        message = str(message) if message is not None else ""
    message = message.strip()

    if not message:
        return "Please enter a valid question."

    if check_emergency(message):
        latency = time.time() - start
        log_request(request_id, message, EMERGENCY_RESPONSE, latency, flagged="emergency_keyword")
        return EMERGENCY_RESPONSE

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history:
        if isinstance(turn, dict) and turn.get("role") in ("user", "assistant") and turn.get("content"):
            messages.append({"role": turn["role"], "content": str(turn["content"])})
    messages.append({"role": "user", "content": message})

    try:
        response = _run_model(messages)
        response = apply_output_guardrails(response)
        latency = time.time() - start
        log_request(request_id, message, response, latency, flagged=None)
        return response
    except Exception as e:
        latency = time.time() - start
        log_request(request_id, message, "", latency, flagged=None, error=str(e))
        logger.error(f"[{request_id}] generation failed: {e}")
        return "Sorry, something went wrong while generating a response. Please try again."


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

with gr.Blocks(title="MediGuide") as demo:
    gr.Markdown(
        """
        # 🩺 MediGuide
        ### QLoRA Fine‑Tuned Medical Conversational Assistant
        MediGuide is based on **Qwen2.5‑1.5B‑Instruct**
        and fine‑tuned using **QLoRA** on medical dialogue data.

        ⚠️ **Disclaimer**
        This demo is for research and educational purposes only.
        It is **not a substitute** for professional medical advice,
        diagnosis, or treatment.
        """
    )

    chatbot = gr.Chatbot(label="MediGuide", height=500)
    message = gr.Textbox(
        label="Your question",
        placeholder="e.g. What are common symptoms of the flu?",
        lines=2,
    )

    with gr.Row():
        send_btn = gr.Button("Send", variant="primary")
        clear_btn = gr.Button("Clear")

    def respond(message, history):
        response = generate_response(message, history)
        history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": response},
        ]
        return "", history

    message.submit(respond, [message, chatbot], [message, chatbot])
    send_btn.click(respond, [message, chatbot], [message, chatbot])
    clear_btn.click(lambda: [], outputs=chatbot)


# ============================================================
# Launch via Gradio SDK (required for ZeroGPU detection),
# then bolt a REST API onto Gradio's own FastAPI app.
# ============================================================

class ChatRequest(BaseModel):
    message: str
    history: list = []


class ChatResponse(BaseModel):
    response: str


demo.queue().launch(
    server_name="0.0.0.0",
    server_port=7860,
    prevent_thread_lock=True,   # let us keep running code after launch() returns
)


@demo.app.get("/health")
def health():
    return {"status": "ok"}


@demo.app.post("/chat", response_model=ChatResponse)
def chat_endpoint(req: ChatRequest):
    response = generate_response(req.message, req.history)
    return ChatResponse(response=response)


demo.block_thread()  # keep the process alive (replaces the old blocking demo.launch())