modules/saathi_chat.py

"""Module 1 — Saathi Chat: general mental-health Q&A + city-based doctor finder.

Flow:
1. User types a mental-health question (in any supported language).
2. Crisis regex runs first. If positive, Claude is bypassed and the helpline banner fires.
3. Otherwise Claude answers as Saathi, in the user's language.
4. User can optionally type a city to get 3-5 pre-curated resources from data/mental_health_resources.json.
5. Student Mode toggle: when ON, shows event-type selector and routes through student coaching prompt.
"""
from __future__ import annotations

from typing import Dict, List

import streamlit as st

from backend.voice_widget import render_mic_button

from backend.claude_client import chat
from backend.i18n import claude_language_name, t
from backend.resources import get_mental_health_resources, list_known_cities
from backend.safeguards import check_crisis, render_crisis_banner
from modules.cognitive_journal import get_cognitive_journal_context



MODULE_NAME = "saathi_chat"
STUDENT_MODULE_NAME = "student_corner"
HISTORY_KEY = "saathi_chat_history"
CITY_KEY = "saathi_chat_city"
MEMORY_TOGGLE_KEY = "saathi_chat_cross_module_memory_enabled"
STUDENT_MODE_KEY = "saathi_chat_student_mode"
STUDENT_EVENT_KEY = "saathi_chat_student_event"

# How many prior messages (not counting the current turn) to resend to Claude.
# 20 = 10 full user↔assistant exchanges. Caps token cost + latency on long demos.
HISTORY_WINDOW = 20

# Event choices for Student Mode (value, i18n key)
STUDENT_EVENT_CHOICES = [
    ("exam", "event_exam"),
    ("placement_interview", "event_placement"),
    ("viva", "event_viva"),
    ("presentation", "event_presentation"),
    ("result_day", "event_result"),
    ("general_burnout", "event_burnout"),
]


def _init_state() -> None:
    if HISTORY_KEY not in st.session_state:
        st.session_state[HISTORY_KEY] = []  # list of {role, content}
    if CITY_KEY not in st.session_state:
        st.session_state[CITY_KEY] = ""
    if MEMORY_TOGGLE_KEY not in st.session_state:
        # Default ON — cross-module awareness is Saathi's differentiator.
        st.session_state[MEMORY_TOGGLE_KEY] = True
    if STUDENT_MODE_KEY not in st.session_state:
        st.session_state[STUDENT_MODE_KEY] = False
    if STUDENT_EVENT_KEY not in st.session_state:
        st.session_state[STUDENT_EVENT_KEY] = "exam"


def _build_cross_module_memory_block() -> str:
    """Return the memory digest iff the user has opted in (default on)."""
    if not st.session_state.get(MEMORY_TOGGLE_KEY, True):
        return ""
    try:
        return get_cognitive_journal_context() or ""
    except Exception:
        # Never let a memory-digest bug break the chat turn.
        return ""




def _render_resources(resources: List[Dict], lang: str) -> None:
    st.markdown(f"### {t('chat_resources_heading', lang)}")
    for r in resources:
        with st.container(border=True):
            st.markdown(f"**{r.get('name', '')}** — *{r.get('type', '')}*")
            if r.get("address"):
                st.markdown(f"📍 {r['address']}")
            if r.get("phone"):
                st.markdown(f"📞 `{r['phone']}`")
            website = r.get("website")
            if website:
                st.markdown(f"🌐 [{website}]({website})")
            specialties = r.get("specialties") or []
            if specialties:
                st.caption("• " + " · ".join(specialties))
            if r.get("cost"):
                st.caption(f"💰 {r['cost']}")
    st.caption(
        "_Resources are curated from Government of India, AIIMS, NIMHANS, and "
        "state mental-health websites. Saathi is not affiliated with any of "
        "these institutions and does not make referrals on their behalf._"
    )


def render(lang: str) -> None:
    """Top-level render function. `lang` is the language code (e.g. 'en', 'hi')."""
    _init_state()

    st.header(t("chat_header", lang))
    st.caption(t("chat_sub", lang))

    # --- Student Mode toggle + Cross-module memory (side by side) ---
    col_student, col_memory = st.columns(2)
    with col_student:
        st.toggle(
            t("chat_student_mode_label", lang),
            key=STUDENT_MODE_KEY,
            help=t("chat_student_mode_help", lang),
        )
    with col_memory:
        with st.expander(t("chat_memory_heading", lang), expanded=False):
            st.toggle(
                t("chat_memory_toggle_label", lang),
                key=MEMORY_TOGGLE_KEY,
                help=t("chat_memory_toggle_help", lang),
            )
            digest_preview = _build_cross_module_memory_block()
            if digest_preview:
                st.caption(t("chat_memory_active_caption", lang))
                st.code(digest_preview, language="markdown")
            else:
                st.caption(t("chat_memory_empty_caption", lang))

    # --- Student Mode: event picker ---
    student_mode_on = st.session_state.get(STUDENT_MODE_KEY, False)
    if student_mode_on:
        st.markdown(f"**{t('student_event_label', lang)}**")
        labels = [t(label_key, lang) for _, label_key in STUDENT_EVENT_CHOICES]
        event_values = [value for value, _ in STUDENT_EVENT_CHOICES]
        current_index = (
            event_values.index(st.session_state[STUDENT_EVENT_KEY])
            if st.session_state[STUDENT_EVENT_KEY] in event_values
            else 0
        )
        chosen_label = st.radio(
            "event",
            options=labels,
            index=current_index,
            horizontal=True,
            label_visibility="collapsed",
            key="saathi_chat_student_event_radio",
        )
        st.session_state[STUDENT_EVENT_KEY] = event_values[labels.index(chosen_label)]

    # --- Conversation history ---
    for msg in st.session_state[HISTORY_KEY]:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    # --- Voice mic (auto-transcribes into chat input) ---
    render_mic_button(lang, target_aria_label=t("chat_input_placeholder", lang))

    # --- Chat input ---
    user_text = st.chat_input(t("chat_input_placeholder", lang), key="saathi_chat_input")
    if user_text:
        # Crisis check BEFORE Claude is invoked
        if check_crisis(user_text):
            st.session_state[HISTORY_KEY].append({"role": "user", "content": user_text})
            with st.chat_message("user"):
                st.markdown(user_text)
            with st.chat_message("assistant"):
                render_crisis_banner(lang)
            return

        # Normal flow → Claude
        st.session_state[HISTORY_KEY].append({"role": "user", "content": user_text})
        with st.chat_message("user"):
            st.markdown(user_text)

        with st.chat_message("assistant"):
            with st.spinner("…"):
                try:
                    # Route through student prompt when Student Mode is ON
                    if student_mode_on:
                        augmented_text = (
                            f"event_type: {st.session_state[STUDENT_EVENT_KEY]}\n"
                            f"situation: {user_text}"
                        )
                        reply = chat(
                            module=STUDENT_MODULE_NAME,
                            user_text=augmented_text,
                            language_name=claude_language_name(lang),
                            max_tokens=2400,
                            extra_context={
                                "cross_module_memory": _build_cross_module_memory_block()
                            },
                        )
                    else:
                        prior_history = st.session_state[HISTORY_KEY][:-1][-HISTORY_WINDOW:]
                        reply = chat(
                            module=MODULE_NAME,
                            user_text=user_text,
                            language_name=claude_language_name(lang),
                            history=prior_history,
                            extra_context={
                                "cross_module_memory": _build_cross_module_memory_block()
                            },
                        )
                except Exception as e:
                    reply = (
                        "I'm having trouble reaching my language model right now. "
                        "If this is urgent, please call Tele-MANAS on **14416** "
                        "(or 1800-89-14416) — Government of India, 24×7, free, 20+ Indian languages. "
                        "For an emergency, call **112**.\n\n"
                        f"_(Technical detail: {e})_"
                    )
            st.markdown(reply)
            st.session_state[HISTORY_KEY].append({"role": "assistant", "content": reply})

    # --- City-based resource finder (always visible under the chat) ---
    st.divider()
    st.markdown(f"**{t('chat_city_prompt', lang)}**")
    col_input, col_button = st.columns([3, 1])
    with col_input:
        city = st.text_input(
            "city",
            value=st.session_state[CITY_KEY],
            placeholder=t("chat_city_placeholder", lang),
            label_visibility="collapsed",
            key="saathi_chat_city_input",
        )
    with col_button:
        find_clicked = st.button(
            t("chat_city_button", lang),
            key="saathi_chat_city_button",
            use_container_width=True,
        )

    if find_clicked and city:
        st.session_state[CITY_KEY] = city
        resources = get_mental_health_resources(city)
        if resources:
            _render_resources(resources, lang)
        else:
            st.warning(t("chat_no_resources", lang))
            known = list_known_cities()
            if known:
                st.caption("Cities I currently know: " + ", ".join(known))