"""
Professional Nurse Advocate Assistant
Powered by Claude claude-opus-4-6 | A-EQUIP Model | NHS OGL v3.0

BYOK — bring your own Anthropic API key. Free forever on Hugging Face.
"""

import streamlit as st
from datetime import date

from pna.rag import PNAKnowledgeBase
from pna.claude_client import stream_response, generate_supervision_note
from pna.export import build_supervision_docx, build_cpd_docx, HAS_DOCX

# ─── Page config ─────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="PNA Assistant | Professional Nurse Advocate",
    page_icon="👨🏾‍⚕️",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ─── Custom CSS ──────────────────────────────────────────────────────────────

st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Fraunces:wght@600;700&family=Inter:wght@400;500;600&display=swap');

html, body, [class*="css"] { font-family: 'Inter', sans-serif; }

.hero-header {
    background: linear-gradient(135deg, #1a2460 0%, #2d3da0 60%, #0d9488 100%);
    padding: 2rem 2rem 1.5rem;
    border-radius: 12px;
    margin-bottom: 1.5rem;
    color: white;
}
.hero-header h1 {
    font-family: 'Fraunces', serif;
    font-size: 2rem;
    margin: 0 0 0.25rem 0;
    color: white;
}
.hero-header p { margin: 0; opacity: 0.85; font-size: 1rem; }

.pill {
    display: inline-block;
    background: rgba(255,255,255,0.15);
    border: 1px solid rgba(255,255,255,0.25);
    border-radius: 99px;
    padding: 0.2rem 0.75rem;
    font-size: 0.8rem;
    margin-top: 0.75rem;
    margin-right: 0.25rem;
}

.disclaimer {
    background: #f8fafc;
    border: 1px solid #e2e8f0;
    border-radius: 8px;
    padding: 0.75rem 1rem;
    font-size: 0.8rem;
    color: #64748b;
    margin-top: 1rem;
}

.ogl-notice {
    font-size: 0.7rem;
    color: #94a3b8;
    margin-top: 0.5rem;
}

[data-testid="stChatMessage"] { border-radius: 12px; }
</style>
""", unsafe_allow_html=True)

# ─── Session state ────────────────────────────────────────────────────────────

defaults = {
    "messages": [],
    "api_key": "",
    "nurse_name": "",
    "session_started": None,
}
for k, v in defaults.items():
    if k not in st.session_state:
        st.session_state[k] = v

# ─── Load RAG knowledge base ──────────────────────────────────────────────────

@st.cache_resource(show_spinner="Loading A-EQUIP knowledge base…")
def load_knowledge_base():
    return PNAKnowledgeBase()

kb = load_knowledge_base()

# ─── Sidebar ─────────────────────────────────────────────────────────────────

with st.sidebar:
    st.markdown("### 👨🏾‍⚕️ PNA Assistant")
    st.caption("A-EQUIP Model · Restorative Supervision")
    st.divider()

    # ── API Key (BYOK) ────────────────────────────────────────────────────────
    st.markdown("#### 🔑 Your Anthropic API Key")
    api_key_input = st.text_input(
        "Enter key to start chatting",
        type="password",
        value=st.session_state.api_key,
        placeholder="sk-ant-api03-...",
        help="Used only for this session and never stored.",
    )
    if api_key_input:
        st.session_state.api_key = api_key_input
        st.success("✅ Key entered")

    st.caption(
        "[Get a free API key →](https://console.anthropic.com) "
        "· Costs ~£0.01 per conversation"
    )

    st.divider()

    # ── Your details ──────────────────────────────────────────────────────────
    st.markdown("#### 📋 Your Details")
    st.caption("Optional — used in supervision note exports.")
    st.session_state.nurse_name = st.text_input(
        "Your name",
        value=st.session_state.nurse_name,
        placeholder="e.g. Nurse Jane Smith",
    )
    cpd_hours = st.number_input(
        "Session length (hours)",
        min_value=0.5, max_value=8.0,
        value=1.0, step=0.5,
    )

    st.divider()

    # ── Session controls ──────────────────────────────────────────────────────
    st.markdown("#### 🔄 Session")
    col1, col2 = st.columns(2)
    with col1:
        if st.button("🗑️ Clear", use_container_width=True):
            st.session_state.messages = []
            st.session_state.session_started = None
            st.rerun()
    with col2:
        if st.button("🆕 New", use_container_width=True):
            st.session_state.messages = []
            st.session_state.session_started = date.today().isoformat()
            st.rerun()

    # ── Exports ───────────────────────────────────────────────────────────────
    if st.session_state.messages:
        st.divider()
        st.markdown("#### 📥 Export")

        if st.button("📄 Generate supervision note", use_container_width=True):
            if not st.session_state.api_key:
                st.error("Please enter your API key first.")
            else:
                with st.spinner("Generating note with Claude…"):
                    note = generate_supervision_note(
                        st.session_state.api_key,
                        st.session_state.messages,
                    )
                    st.session_state["last_note"] = note

        if "last_note" in st.session_state:
            note = st.session_state["last_note"]
            st.text_area("Preview", note, height=200)

            if HAS_DOCX:
                docx_bytes = build_supervision_docx(note, st.session_state.nurse_name)
                if docx_bytes:
                    st.download_button(
                        "⬇️ Supervision note (.docx)",
                        data=docx_bytes,
                        file_name=f"PNA_Supervision_{date.today()}.docx",
                        mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                        use_container_width=True,
                    )

                cpd_bytes = build_cpd_docx(note, st.session_state.nurse_name, cpd_hours)
                if cpd_bytes:
                    st.download_button(
                        "⬇️ NMC CPD Record (.docx)",
                        data=cpd_bytes,
                        file_name=f"NMC_CPD_{date.today()}.docx",
                        mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                        use_container_width=True,
                    )
            else:
                st.download_button(
                    "⬇️ Download as .txt",
                    data=note,
                    file_name=f"PNA_Supervision_{date.today()}.txt",
                    mime="text/plain",
                    use_container_width=True,
                )

    st.divider()
    st.markdown(
        '<p class="ogl-notice">Contains public sector information licensed under the '
        '<a href="https://www.nationalarchives.gov.uk/doc/open-government-licence/version/3/" '
        'target="_blank">Open Government Licence v3.0</a> — NHS England.</p>',
        unsafe_allow_html=True,
    )

# ─── Main content ─────────────────────────────────────────────────────────────

st.markdown("""
<div class="hero-header">
    <h1>👨🏾‍⚕️ Professional Nurse Advocate Assistant</h1>
    <p>Your AI guide to the A-EQUIP model, restorative supervision, and quality improvement</p>
    <span class="pill">🇬🇧 NHS England · A-EQUIP Model</span>
    <span class="pill">NMC Standards</span>
    <span class="pill">OGL v3.0</span>
</div>
""", unsafe_allow_html=True)

# ── Welcome / quick start ─────────────────────────────────────────────────────

if not st.session_state.messages:
    col1, col2, col3 = st.columns(3)
    with col1:
        st.markdown("""
**🌿 Restorative Supervision**
Start a reflective conversation to process the emotional impact of your clinical work.
""")
    with col2:
        st.markdown("""
**📚 A-EQUIP Learning**
Ask about any of the four A-EQUIP functions or explore the PNA role in depth.
""")
    with col3:
        st.markdown("""
**✅ PAQI Planning**
Get support developing your Personal Action for Quality Improvement.
""")

    st.markdown("#### 💬 Try asking…")
    examples = [
        "What is the A-EQUIP model and why does it matter for nurses?",
        "I've had a really difficult week on the ward. I don't know where to start.",
        "How do I facilitate a restorative supervision session for the first time?",
        "What does a PNA actually do day-to-day?",
        "Help me reflect on a challenging patient interaction using Gibbs cycle.",
        "What's the difference between the normative and formative functions?",
    ]
    cols = st.columns(2)
    for i, ex in enumerate(examples):
        with cols[i % 2]:
            if st.button(ex, key=f"ex_{i}", use_container_width=True):
                st.session_state.messages.append({"role": "user", "content": ex})
                st.rerun()

# ── Chat history ──────────────────────────────────────────────────────────────

for msg in st.session_state.messages:
    with st.chat_message(msg["role"], avatar="👩‍⚕️" if msg["role"] == "user" else "👨🏾‍⚕️"):
        st.markdown(msg["content"])

# ── Chat input ────────────────────────────────────────────────────────────────

prompt = st.chat_input(
    "Ask me about A-EQUIP, restorative supervision, or the PNA role…",
    disabled=not st.session_state.api_key,
)

if not st.session_state.api_key:
    st.info(
        "👈 Enter your Anthropic API key in the sidebar to start chatting. "
        "[Get a free key →](https://console.anthropic.com)"
    )

if prompt and st.session_state.api_key:
    if st.session_state.session_started is None:
        st.session_state.session_started = date.today().isoformat()

    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar="👩‍⚕️"):
        st.markdown(prompt)

    context = kb.search(prompt) if kb.index is not None else ""

    with st.chat_message("assistant", avatar="👨🏾‍⚕️"):
        placeholder = st.empty()
        full_response = ""

        try:
            for chunk in stream_response(
                api_key=st.session_state.api_key,
                history=st.session_state.messages[:-1],
                user_message=prompt,
                context=context,
            ):
                full_response += chunk
                placeholder.markdown(full_response + "▋")

            placeholder.markdown(full_response)

        except Exception as e:
            err = str(e)
            if "authentication" in err.lower() or "api_key" in err.lower():
                full_response = "❌ Invalid API key. Please check your key in the sidebar."
            elif "rate" in err.lower():
                full_response = "⏳ Rate limit reached. Please wait a moment and try again."
            else:
                full_response = f"❌ An error occurred: {err}"
            placeholder.error(full_response)

    st.session_state.messages.append({"role": "assistant", "content": full_response})

# ─── Footer ───────────────────────────────────────────────────────────────────

st.divider()
st.markdown("""
<div class="disclaimer">
⚠️ <strong>Clinical Disclaimer:</strong> This tool is for educational purposes only.
It does not provide clinical advice, diagnosis, or treatment recommendations.
It is not a replacement for human supervision or your employer's formal clinical governance processes.
If you are experiencing a clinical emergency or safeguarding concern, follow your Trust's protocols immediately.
<br/><br/>
<span class="ogl-notice">
Contains public sector information licensed under the
<a href="https://www.nationalarchives.gov.uk/doc/open-government-licence/version/3/" target="_blank">Open Government Licence v3.0</a>
— NHS England.
Built by <a href="https://nursingcitizendevelopment.com" target="_blank">Lincoln Gombedza</a> · CQAI ·
<a href="https://github.com/Clinical-Quality-Artifical-Intelligence/Professional-Nurse-Advocate-Assistant" target="_blank">Open Source</a>
</span>
</div>
""", unsafe_allow_html=True)