SYSTEM_PROMPT = """You are a medical AI assistant. You analyze medical images
and answer questions using retrieved reference material. Be concise and
clinically precise. Always note limitations. Do not provide diagnoses —
provide information to support a clinician's review. Provide structured and bulleted answers when required.
DO NOT make up references or hallucinate information. If you are unsure, say so.
"""

def build_medgemma_prompt(question, history, rag_chunks, has_image):
    # Bound history to last 4 turns — first line of defense against KV OOM
    history = history[-8:]  # 4 user + 4 assistant

    if rag_chunks:
        rag_block = "\n\n".join(
            f"[Source: {c['source']}]\n{c['text']}" for c in rag_chunks
        )
        reference_section = f"Reference material:\n{rag_block}\n\n"
    else:
        # Don't force in irrelevant chunks just to have something -- an empty
        # or mismatched reference block burns reasoning tokens on the model
        # trying to reconcile material that has nothing to do with the
        # question, which can eat the whole response budget before it ever
        # gets to an actual answer.
        reference_section = (
            "No relevant reference material was found for this question — "
            "answer from your own knowledge instead.\n\n"
        )

    # Gemma's chat template doesn't properly support a separate "system" role
    # (especially combined with a multimodal user turn right after it), so the
    # instructions are folded into the user turn's text instead.
    messages = list(history)

    user_content = (
        f"{SYSTEM_PROMPT}\n"
        f"{reference_section}"
        f"Question: {question}"
    )
    if has_image:
        user_content = "Analyze the attached medical image.\n\n" + user_content

    messages.append({"role": "user", "content": user_content})
    return messages