import os
import re
import streamlit as st
from openai import OpenAI

# Configure page style
st.set_page_config(page_title="Clinical ReAct Agent — MIS 752", layout="wide")
st.title("🩺 Clinical ReAct Agent & Supervisor Dashboard")
st.markdown(
    "**MIS 752 · Lab 9: Build the Paperwork Machine** — Custom ReAct loop (`Thought → Action → Observation`) "
    "with `stop=['Observation:']` hallucination prevention and automated human-in-the-loop escalation policies."
)

# Sidebar - Secrets & Settings
st.sidebar.header("⚙️ Agent & Supervisor Settings")

default_key = os.environ.get("OPENROUTER_API_KEY") or os.environ.get("OPENAI_API_KEY") or ""
api_key = st.sidebar.text_input(
    "API Key (OpenRouter or OpenAI)",
    value=default_key,
    type="password",
    placeholder="sk-or-v1-... or sk-..."
)

is_openrouter = api_key.strip().startswith("sk-or-")
provider_label = "OpenRouter (https://openrouter.ai/api/v1)" if is_openrouter else "OpenAI Default"
st.sidebar.caption(f"🔌 Active Endpoint: **{provider_label}**")

if is_openrouter:
    model_options = [
        "openai/gpt-4o-mini",
        "qwen/qwen3-8b",
        "meta-llama/llama-3.3-70b-instruct",
        "anthropic/claude-3.5-haiku"
    ]
else:
    model_options = ["gpt-4o-mini", "gpt-4o", "gpt-3.5-turbo"]

model_name = st.sidebar.selectbox("LLM Model Name", model_options)
max_steps = st.sidebar.slider("Step Budget (Max Steps)", 2, 8, 6)
supervision_policy = st.sidebar.selectbox(
    "Supervision Policy (Step 9.2)",
    ["flagged", "every", "audit"],
    help="Controls when the human supervisor is interrupted."
)

# Clinic data structures
REFERENCE = {
    "metformin":     "Contraindicated when eGFR is below 30. Between 30 and 45, do not start it; if the patient is already taking it, halve the maximum to 1000 mg/day and recheck eGFR every 3 months.",
    "a1c target":    "General adult target is 7.0 percent. For an older or frail patient, 8.0 percent is reasonable: a tighter target buys little and costs hypoglycaemia risk.",
    "ckd stage":     "G1-G2 is eGFR 60 or above. G3a is 45-59. G3b is 30-44. G4 is 15-29. G5 is below 15, which is kidney failure.",
    "alert fatigue": "Degraded human response to a high volume of alarms, caused by the alarms themselves. Around 90 percent of drug-interaction alerts are overridden.",
    "lisinopril":    "An ACE inhibitor. Monitor potassium and renal function. Hold it if the patient becomes acutely unwell or dehydrated.",
}

CHARTS = {
    "alvarez": "Alvarez, F, 68. Type 2 diabetes. eGFR 28. A1c 9.1 percent. Current: metformin 1000 mg twice daily, lisinopril 10 mg daily.",
    "okafor":  "Okafor, M, 54. Hypertension. eGFR 71. A1c 5.8 percent. Current: lisinopril 20 mg daily.",
    "bianchi": "Bianchi, F, 81. Type 2 diabetes, frail. eGFR 39. A1c 8.4 percent. Current: metformin 500 mg twice daily.",
}

with st.sidebar.expander("📚 View Reference & Patient Charts"):
    st.markdown("**Clinic Reference Topics:** `" + "`, `".join(REFERENCE.keys()) + "`")
    st.markdown("**Patient Charts:** `" + "`, `".join(CHARTS.keys()) + "`")

def lookup(topic: str) -> str:
    hit = REFERENCE.get(topic.strip().lower())
    return hit if hit else f"'{topic}' is not in the clinic reference."

def get_patient(name: str) -> str:
    hit = CHARTS.get(name.strip().lower())
    return hit if hit else f"There is no patient called '{name}' on this list."

TOOLS = {"lookup": lookup, "get_patient": get_patient}

def run_tool(name, arg):
    fn = TOOLS.get(name)
    if fn is None:
        return f"There is no tool called '{name}'. Available: {', '.join(TOOLS)}.", False
    try:
        return fn(arg), True
    except Exception as e:
        return f"The tool failed: {type(e).__name__}: {e}", False

def parse_action(text: str):
    tool = re.search(r"Action:\s*([a-zA-Z_]+)", text)
    arg  = re.search(r"Action Input:\s*(.+)", text)
    if not tool:
        return None, None
    return tool.group(1).strip(), (arg.group(1).strip() if arg else "")

def escalate_step(stage: dict, policy: str = "flagged") -> str:
    where = f"step {stage.get('step', '?')} ({stage.get('stage', '?')})"
    tool = stage.get("tool") or "no tool"
    ok = bool(stage.get("ok", True))

    if policy == "every":
        return f"INTERRUPT - {where}, {tool}: policy is review-everything."
    elif policy == "audit":
        return f"LOG - {where}, {tool}: recorded for later review, nobody paged."
    elif policy == "flagged":
        if not ok:
            return f"INTERRUPT - {where}, {tool}: the agent reported a failure here."
        return f"SILENT - {where}, {tool}: nothing flagged."
    else:
        return f"INTERRUPT - unknown policy '{policy}'. Defaulting to interrupt."

SYSTEM_PROMPT = '''You answer clinical questions by looking things up. You have two tools:

lookup: look up one topic in the clinic reference. Input: a short topic, e.g. "metformin".
get_patient: look up a patient chart by surname. Input: a surname, e.g. "Alvarez".

ALWAYS answer in this exact format, ONE step at a time:

Thought: what you know so far and what you need next
Action: lookup
Action Input: the topic to look up

Then STOP and wait. You will be given an Observation with the result.
Never write the Observation yourself.

When you have enough to answer, write instead:

Thought: I now know the final answer
Final Answer: your answer, in one or two sentences
/no_think'''

preset = st.selectbox(
    "Choose a preset lab question or customize below:",
    [
        "Is Alvarez still OK on her metformin? Check her chart first.",
        "What share of drug-interaction alerts get overridden, and why does that matter for a unit that fires 240 alerts a shift?",
        "What does our hospital night shift override policy say about 240 alerts in one night?",
        "Check Bianchi's chart and our a1c target and metformin rules. Does her regimen need adjustment?"
    ]
)

query = st.text_input("Ask the clinical agent a question:", value=preset)

if st.button("▶️ Run Agent", type="primary"):
    if not api_key:
        st.error("Please enter an OpenRouter (`sk-or-v1-...`) or OpenAI (`sk-...`) API Key in the sidebar to run the agent.")
    else:
        try:
            base_url = "https://openrouter.ai/api/v1" if is_openrouter else None
            client = OpenAI(api_key=api_key.strip(), base_url=base_url)
            messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": query}]
            converged = False

            st.subheader("🧠 Agent Execution Trace")

            for step in range(1, max_steps + 1):
                with st.spinner(f"Thinking (Step {step})..."):
                    r = client.chat.completions.create(
                        model=model_name,
                        messages=messages,
                        temperature=0,
                        max_tokens=1000,
                        stop=["Observation:"]
                    )
                    reply = r.choices[0].message.content or ""
                    reply = re.sub(r"<think>.*?</think>", "", reply, flags=re.S)
                    reply = re.sub(r"<think>.*$", "", reply, flags=re.S).strip()

                model_stage = {"step": step, "stage": "model", "text": reply, "tool": None, "ok": bool(reply)}

                with st.chat_message("assistant"):
                    st.markdown(f"**Step {step} Reasoning** · `{escalate_step(model_stage, supervision_policy)}`")
                    st.text(reply or "[empty reply]")

                if "Final Answer:" in reply:
                    answer = reply.split("Final Answer:", 1)[1].strip()
                    st.success(f"✅ **Final Answer (Converged in {step} steps):** {answer}")
                    converged = True
                    break

                tool, arg = parse_action(reply)
                if tool is None:
                    obs, ok = "Could not parse action format. Please follow Action/Action Input syntax.", False
                else:
                    obs, ok = run_tool(tool, arg)

                obs_stage = {"step": step, "stage": "observation", "text": obs, "tool": tool, "ok": ok}

                with st.chat_message("user"):
                    st.markdown(f"**Step {step} Observation (Tool: `{tool}`)** · `{escalate_step(obs_stage, supervision_policy)}`")
                    if ok:
                        st.info(obs)
                    else:
                        st.error(obs)

                messages.append({"role": "assistant", "content": reply})
                messages.append({"role": "user", "content": f"Observation: {obs}"})

            if not converged:
                st.warning(f"⚠️ Step budget exhausted (`max_steps={max_steps}`) before reaching a Final Answer.")

        except Exception as e:
            st.error(f"An error occurred: {e}")