"""
P3 Myeloma - AI Follow-up Chatbot
==================================
Takes a participant's answers to the 16 Shared Decision Making questions
and runs an adaptive follow-up conversation using Google's Gemini API.

Deploy on Hugging Face Spaces (SDK: Gradio).
Set your Gemini key as a Space Secret named GEMINI_API_KEY,
or paste it in the UI field at runtime.
"""

import os
import json
from google import genai
from google.genai import types
import gradio as gr

# ----------------------------------------------------------------------
# The 16 profile questions (verbatim from the intake form)
# ----------------------------------------------------------------------
QUESTIONS = [
    "I prefer my healthcare team (hematologist, nurse practitioner, or physician assistant), and I collaborate in deciding which treatment for relapsed or refractory myeloma is best for me",
    "It is important for me to understand my treatment options for relapsed or refractory myeloma",
    "I trust my healthcare team very much, that's why I leave it up to them to recommend therapy that is right for me",
    "I live alone with no caregiver",
    "I have strong social support and a network that can help with my treatment appointments",
    "I will do whatever it takes to so I can be cured, cancer-free, kill all the myeloma, or at least be in remission and live a long life",
    "I would like to take an aggressive approach to treat my myeloma",
    "I am willing to endure as many side effects as possible to control my myeloma",
    "I prefer to receive treatment in an outpatient setting",
    "I prefer to take medications at home",
    "I prefer to take the least possible amount of pills to control my cancer",
    "Quality of life is more important to me than quantity of life",
    "Clinical drug trial participation is of interest to me",
    "My out-of-pocket cost of treatment is important to me",
    "I prefer to continue an active lifestyle during my myeloma treatment",
    "I worry about how my treatment will affect future treatment options",
]

# Short tags used for the summary view
QUESTION_TAGS = [
    "Shared decision-making",
    "Wants to understand options",
    "Defers to care team",
    "Lives alone / no caregiver",
    "Strong social support",
    "Goal: cure / long life",
    "Wants aggressive approach",
    "Tolerant of side effects",
    "Prefers outpatient",
    "Prefers meds at home",
    "Prefers fewer pills",
    "Quality over quantity",
    "Interested in clinical trials",
    "Out-of-pocket cost matters",
    "Wants active lifestyle",
    "Worries about future options",
]

MODEL = "gemini-2.5-flash"
MAX_FOLLOWUPS = 8  # cap the adaptive conversation

# ----------------------------------------------------------------------
# Prompt building
# ----------------------------------------------------------------------
def build_profile_text(answers):
    """answers: list of 'Yes'/'No' aligned to QUESTIONS."""
    lines = []
    for q, tag, a in zip(QUESTIONS, QUESTION_TAGS, answers):
        lines.append(f"- [{a.upper()}] ({tag}) {q}")
    return "\n".join(lines)


def detect_tensions(answers):
    """Flag notable patterns/contradictions worth probing. Pure logic, no API."""
    a = {i: (answers[i].strip().lower() == "yes") for i in range(len(answers))}
    flags = []

    # Aggressive cure-seeking vs quality-over-quantity
    if (a.get(5) or a.get(6) or a.get(7)) and a.get(11):
        flags.append("Wants an aggressive/curative approach but also values quality of life over quantity. Worth clarifying how they weigh these when they conflict.")
    # Willing to endure side effects vs wants active lifestyle / quality of life
    if a.get(7) and (a.get(14) or a.get(11)):
        flags.append("Willing to endure many side effects, yet wants to stay active / prioritizes quality of life. Probe acceptable side-effect threshold.")
    # Lives alone AND no strong social support -> logistics risk
    if a.get(3) and not a.get(4):
        flags.append("Lives alone with no caregiver and limited social support. Treatment logistics and safety monitoring need exploration.")
    # Lives alone but prefers home meds -> safety of self-administration
    if a.get(3) and a.get(9):
        flags.append("Lives alone but prefers taking medications at home. Explore support for safe self-administration and side-effect monitoring.")
    # Defers fully to team but also wants to collaborate / understand options
    if a.get(2) and (a.get(0) or a.get(1)):
        flags.append("Says they leave decisions to the care team, yet also wants to collaborate / understand options. Clarify how involved they actually want to be.")
    # Cost matters + trial interest (trials can offset cost) — opportunity, not conflict
    if a.get(13) and a.get(12):
        flags.append("Cost matters and they're open to clinical trials — trials may reduce drug cost; worth surfacing.")
    # Fewer pills preference but aggressive
    if a.get(10) and (a.get(6) or a.get(7)):
        flags.append("Prefers the fewest pills possible but wants an aggressive approach. Explore tolerance for treatment intensity vs convenience.")
    # Outpatient + lives alone
    if a.get(8) and a.get(3):
        flags.append("Prefers outpatient treatment but lives alone — explore transport and post-visit support.")
    return flags


def system_instruction():
    return (
        "You are a warm, plain-spoken health navigator helping a multiple myeloma patient "
        "(relapsed or refractory) prepare for a shared decision-making conversation with their "
        "care team. You are NOT a doctor and you never give medical advice, diagnoses, dosing, or "
        "treatment recommendations. Your job is ONLY to ask thoughtful follow-up questions that help "
        "the patient clarify their own values, priorities, constraints, and concerns.\n\n"
        "RULES:\n"
        "1. Ask exactly ONE question per turn. Keep it to 1-3 sentences, conversational, jargon-free.\n"
        "2. Base each question on the patient's profile and their previous answers. Prioritize the "
        "FLAGGED TENSIONS provided — gently explore apparent contradictions without judgment.\n"
        "3. Do not repeat questions already asked. Build on what they say.\n"
        "4. Never recommend or rank treatments. If the patient asks for medical advice, kindly "
        "redirect them to discuss it with their care team, then continue with a follow-up question.\n"
        "5. Tone: empathetic, respectful, never alarming.\n"
        "6. When you judge that you have enough to summarize their priorities (or after several "
        "exchanges), instead of asking another question, output a final summary. Begin that final "
        "message with the exact token <<SUMMARY>> on its own line, then 4-7 short bullet points "
        "capturing the patient's key priorities, constraints, and questions to raise with their care team."
    )


# ----------------------------------------------------------------------
# Gemini client helpers
# ----------------------------------------------------------------------
def get_client(user_key):
    key = (user_key or "").strip() or os.environ.get("GEMINI_API_KEY", "").strip()
    if not key:
        return None, "No API key found. Paste your Gemini API key above, or set GEMINI_API_KEY as a Space secret."
    try:
        client = genai.Client(api_key=key)
        return client, None
    except Exception as e:
        return None, f"Could not initialize Gemini client: {e}"


def to_gemini_contents(history):
    """history: list of {'role': 'user'|'assistant', 'content': str}
    Returns Gemini Content list (assistant -> 'model')."""
    contents = []
    for m in history:
        role = "model" if m["role"] == "assistant" else "user"
        contents.append(types.Content(role=role, parts=[types.Part(text=m["content"])]))
    return contents


def gemini_reply(client, convo, n_followups):
    """convo: internal history list. Returns assistant text."""
    contents = to_gemini_contents(convo)
    # Nudge toward summarizing as we approach the cap
    extra = ""
    if n_followups >= MAX_FOLLOWUPS - 1:
        extra = "\n\nYou have asked enough questions. Provide the final <<SUMMARY>> now."
    cfg = types.GenerateContentConfig(
        system_instruction=system_instruction() + extra,
        temperature=0.7,
        max_output_tokens=600,
    )
    resp = client.models.generate_content(model=MODEL, contents=contents, config=cfg)
    return (resp.text or "").strip()


# ----------------------------------------------------------------------
# Gradio app
# ----------------------------------------------------------------------
def start_chat(api_key, *radio_values):
    answers = [v if v in ("Yes", "No") else None for v in radio_values]
    if any(v is None for v in answers):
        missing = [i + 1 for i, v in enumerate(answers) if v is None]
        gr.Warning(f"Please answer all 16 questions. Missing: {missing}")
        return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))

    client, err = get_client(api_key)
    if err:
        gr.Warning(err)
        return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))

    profile = build_profile_text(answers)
    flags = detect_tensions(answers)
    flag_text = "\n".join(f"- {f}" for f in flags) if flags else "- (No obvious contradictions; explore their highest-stakes priorities.)"

    kickoff = (
        "Here is the patient's completed profile (16 yes/no answers):\n\n"
        f"{profile}\n\n"
        "FLAGGED TENSIONS / PRIORITIES TO EXPLORE:\n"
        f"{flag_text}\n\n"
        "Begin the follow-up conversation. Warmly acknowledge in one short sentence, then ask your "
        "FIRST single follow-up question."
    )

    convo = [{"role": "user", "content": kickoff}]
    try:
        reply = gemini_reply(client, convo, 0)
    except Exception as e:
        gr.Warning(f"Gemini error: {e}")
        return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))

    convo.append({"role": "assistant", "content": reply})
    chat_display = [{"role": "assistant", "content": reply}]
    return (
        chat_display,          # chatbot
        chat_display,          # display_state
        convo,                 # convo_state
        1,                     # followup count
        gr.update(visible=False),  # hide intake
        gr.update(visible=True),   # show chat
    )


def respond(user_msg, chat_display, convo, n_followups, api_key):
    user_msg = (user_msg or "").strip()
    if not user_msg:
        return chat_display, chat_display, convo, n_followups, ""

    client, err = get_client(api_key)
    if err:
        gr.Warning(err)
        return chat_display, chat_display, convo, n_followups, user_msg

    convo = convo + [{"role": "user", "content": user_msg}]
    chat_display = chat_display + [{"role": "user", "content": user_msg}]

    try:
        reply = gemini_reply(client, convo, n_followups)
    except Exception as e:
        gr.Warning(f"Gemini error: {e}")
        return chat_display, chat_display, convo, n_followups, ""

    convo = convo + [{"role": "assistant", "content": reply}]

    display_reply = reply
    if "<<SUMMARY>>" in reply:
        display_reply = reply.replace("<<SUMMARY>>", "").strip()
        display_reply = "**Your Priorities Summary**\n\n" + display_reply + (
            "\n\n_This summary reflects what you shared. Please bring it to your care team. "
            "It is not medical advice._"
        )

    chat_display = chat_display + [{"role": "assistant", "content": display_reply}]
    return chat_display, chat_display, convo, n_followups + 1, ""


def reset_all():
    radio_resets = [gr.update(value=None) for _ in QUESTIONS]
    return (
        [],   # chatbot
        [],   # display_state
        [],   # convo_state
        0,    # followups
        gr.update(visible=True),   # intake
        gr.update(visible=False),  # chat
        *radio_resets,
    )


CSS = """
.q-card {border:1px solid #e5e5e5; border-radius:12px; padding:14px 16px; margin-bottom:10px; background:#fff;}
#title {color:#7a0c2e; font-weight:700;}
footer {visibility:hidden;}
"""

with gr.Blocks(title="P3 Myeloma Follow-up Chatbot", css=CSS,
               theme=gr.themes.Soft(primary_hue="red")) as demo:
    gr.Markdown("# <span id='title'>AI-Enhanced Personalized Patient Profile (P3) Myeloma</span>")
    gr.Markdown("### Follow-up Questions for Shared Decision Making")

    with gr.Accordion("Gemini API Key (optional if set as Space secret GEMINI_API_KEY)", open=False):
        api_key = gr.Textbox(label="GEMINI_API_KEY", type="password",
                             placeholder="Paste your key here, or leave blank to use the Space secret")

    convo_state = gr.State([])
    display_state = gr.State([])
    followups = gr.State(0)

    # ---- Intake panel ----
    with gr.Group(visible=True) as intake_panel:
        gr.Markdown("**Answer all 16 questions, then click Start Follow-up.**")
        radios = []
        for i, q in enumerate(QUESTIONS):
            with gr.Row(elem_classes="q-card"):
                r = gr.Radio(["Yes", "No"], label=f"{i+1}. {q}")
                radios.append(r)
        start_btn = gr.Button("Start Follow-up", variant="primary")

    # ---- Chat panel ----
    with gr.Group(visible=False) as chat_panel:
        chatbot = gr.Chatbot(label="Follow-up Conversation", type="messages", height=460)
        with gr.Row():
            msg = gr.Textbox(placeholder="Type your answer...", show_label=False, scale=8)
            send_btn = gr.Button("Send", variant="primary", scale=1)
        restart_btn = gr.Button("Start Over")

    start_btn.click(
        start_chat,
        inputs=[api_key] + radios,
        outputs=[chatbot, display_state, convo_state, followups, intake_panel, chat_panel],
    )

    send_btn.click(
        respond,
        inputs=[msg, display_state, convo_state, followups, api_key],
        outputs=[chatbot, display_state, convo_state, followups, msg],
    )
    msg.submit(
        respond,
        inputs=[msg, display_state, convo_state, followups, api_key],
        outputs=[chatbot, display_state, convo_state, followups, msg],
    )

    restart_btn.click(
        reset_all,
        inputs=None,
        outputs=[chatbot, display_state, convo_state, followups, intake_panel, chat_panel] + radios,
    )

if __name__ == "__main__":
    demo.launch()