"""Saathi — a multilingual mental-health companion for India.

This is the Streamlit entry point. It wires together:
- Language selector (7 Indian languages, Claude is natively multilingual)
- One-time consent acknowledgment gate (ethical alignment)
- Persistent "not a professional" banner across every tab
- 5 feature tabs (Chat, Legal Aid, Thought Diary, My Patterns, Soothe Corner)
- Session-only sidebar conversations, similar to ChatGPT's left rail
- "What happens to my data?" privacy explainer in sidebar
- Footer with privacy disclaimer

Zero persistence: everything lives in st.session_state. Close the browser, it's gone.
"""
from __future__ import annotations

import copy
from datetime import datetime

import streamlit as st

from backend.claude_client import get_active_provider_label
from backend.i18n import DEFAULT_LANGUAGE, LANGUAGES, native_label, t
from modules import (
    cognitive_journal,
    legal_aid,
    my_patterns,
    saathi_chat,
    soothe_poetry,
)

st.set_page_config(
    page_title="Saathi — Mental Health Companion",
    page_icon="🫂",
    layout="wide",
    initial_sidebar_state="expanded",
)


LANGUAGE_KEY = "saathi_language"
SESSIONS_KEY = "saathi_sessions"
ACTIVE_SESSION_KEY = "saathi_active_session"
SESSION_COUNTER_KEY = "saathi_session_counter"
CONSENT_KEY = "saathi_consent_given"

SESSION_STATE_KEYS = [
    # Saathi Chat
    "saathi_chat_history",
    "saathi_chat_city",
    # Legal Aid
    "legal_aid_situation",
    "legal_aid_category",
    "legal_aid_sections",
    "legal_aid_response",
    "legal_aid_letter",
    # Student Mode (now in Chat)
    "saathi_chat_student_mode",
    "saathi_chat_student_event",
    # My Patterns
    "my_patterns_songs",
    "my_patterns_activities",
    # Thought Diary
    "cognitive_journal_entries",
    "cognitive_journal_last",
    "cognitive_journal_section",
    "cognitive_journal_widget_version",
    "cognitive_journal_notice",
    "cognitive_journal_phq9_score",
    "cognitive_journal_phq9_answers",
    "cognitive_journal_phq9_taken_at",
    "cognitive_journal_phq9_history",
    "cognitive_journal_phq9_item9_positive",
    "cognitive_journal_phq9_widget_version",
    "cognitive_journal_gad7_score",
    "cognitive_journal_gad7_answers",
    "cognitive_journal_gad7_taken_at",
    "cognitive_journal_gad7_history",
    "cognitive_journal_gad7_widget_version",
    "cognitive_journal_checkins",
    "cognitive_journal_legacy_cleanup_done",
    # Soothe Corner
    "soothe_feeling",
    "soothe_poem",
    "soothe_songs",

]

SESSION_WIDGET_KEYS_TO_CLEAR = [
    "saathi_chat_city_input",
    "legal_aid_situation_input",
    "saathi_chat_student_event_radio",
    "cognitive_journal_input",
    "cognitive_journal_checkin_mood",
    "cognitive_journal_checkin_sleep",
    "cognitive_journal_checkin_stress",
    "soothe_feeling_input",

]

SESSION_WIDGET_KEY_PREFIXES_TO_CLEAR = [
    "cognitive_journal_input_v",
    "cognitive_journal_phq9_v",
    "cognitive_journal_gad7_v",
]


_CUSTOM_CSS = """
<style>
/* Header wordmark band */
.saathi-header {
    background: linear-gradient(135deg, #6366F1 0%, #8B5CF6 45%, #EC4899 100%);
    padding: 28px 36px 26px 36px;
    border-radius: 18px;
    color: #FFFFFF;
    box-shadow: 0 8px 28px rgba(99, 102, 241, 0.18);
    margin-bottom: 8px;
}
.saathi-header .saathi-wordmark {
    font-family: "Georgia", "Palatino Linotype", serif;
    font-size: 2.7rem;
    font-weight: 700;
    letter-spacing: -0.01em;
    line-height: 1.1;
    margin: 0;
}
.saathi-header .saathi-sub {
    font-family: "Georgia", "Palatino Linotype", serif;
    font-size: 1.05rem;
    opacity: 0.93;
    margin: 6px 0 0 0;
    font-style: italic;
}
.saathi-header .saathi-devanagari {
    font-size: 1.35rem;
    font-weight: 400;
    opacity: 0.85;
    margin-left: 8px;
    letter-spacing: 0;
}
/* Consent screen centre */
.saathi-consent {
    max-width: 620px;
    margin: 60px auto;
    text-align: center;
    padding: 40px 32px;
    border-radius: 18px;
    background: rgba(99, 102, 241, 0.06);
    border: 1px solid rgba(99, 102, 241, 0.15);
}
.saathi-consent h2 {
    color: #4f46e5;
    margin-bottom: 16px;
}
/* Tab row — slightly larger, softer */
[data-testid="stTabs"] [data-baseweb="tab-list"] {
    gap: 4px;
    border-bottom: 1px solid rgba(99, 102, 241, 0.12);
}
[data-testid="stTabs"] [data-baseweb="tab"] {
    padding: 10px 18px;
    font-size: 0.98rem;
    font-weight: 500;
}
[data-testid="stTabs"] [aria-selected="true"] {
    background: rgba(99, 102, 241, 0.08);
    border-radius: 8px 8px 0 0;
}
footer {visibility: hidden;}
</style>
"""


def _init_language() -> None:
    if LANGUAGE_KEY not in st.session_state:
        st.session_state[LANGUAGE_KEY] = DEFAULT_LANGUAGE


def _new_session_title(index: int) -> str:
    return f"Session {index}"


def _init_sessions() -> None:
    if SESSION_COUNTER_KEY not in st.session_state:
        st.session_state[SESSION_COUNTER_KEY] = 1
    if SESSIONS_KEY not in st.session_state:
        st.session_state[SESSIONS_KEY] = {
            "session_1": {
                "title": _new_session_title(1),
                "created_at": datetime.now().isoformat(timespec="minutes"),
                "state": {},
            }
        }
    if ACTIVE_SESSION_KEY not in st.session_state:
        st.session_state[ACTIVE_SESSION_KEY] = "session_1"


def _current_session_id() -> str:
    _init_sessions()
    active_id = st.session_state[ACTIVE_SESSION_KEY]
    sessions = st.session_state[SESSIONS_KEY]
    if active_id not in sessions:
        active_id = next(iter(sessions))
        st.session_state[ACTIVE_SESSION_KEY] = active_id
    return active_id


def _save_active_session() -> None:
    active_id = _current_session_id()
    state = {}
    for key in SESSION_STATE_KEYS:
        if key in st.session_state:
            state[key] = copy.deepcopy(st.session_state[key])
    st.session_state[SESSIONS_KEY][active_id]["state"] = state


def _restore_session(session_id: str) -> None:
    sessions = st.session_state[SESSIONS_KEY]
    if session_id not in sessions:
        return
    restored = sessions[session_id].get("state", {})
    for key in SESSION_STATE_KEYS:
        if key in restored:
            st.session_state[key] = copy.deepcopy(restored[key])
        elif key in st.session_state:
            del st.session_state[key]
    for key in SESSION_WIDGET_KEYS_TO_CLEAR:
        if key in st.session_state:
            del st.session_state[key]
    for key in list(st.session_state.keys()):
        if any(key.startswith(prefix) for prefix in SESSION_WIDGET_KEY_PREFIXES_TO_CLEAR):
            del st.session_state[key]
    st.session_state[ACTIVE_SESSION_KEY] = session_id


def _create_session() -> None:
    _save_active_session()
    st.session_state[SESSION_COUNTER_KEY] += 1
    idx = st.session_state[SESSION_COUNTER_KEY]
    session_id = f"session_{idx}"
    st.session_state[SESSIONS_KEY][session_id] = {
        "title": _new_session_title(idx),
        "created_at": datetime.now().isoformat(timespec="minutes"),
        "state": {},
    }
    _restore_session(session_id)


def _switch_session(session_id: str) -> None:
    if session_id == st.session_state[ACTIVE_SESSION_KEY]:
        return
    _save_active_session()
    _restore_session(session_id)


def _render_session_sidebar(lang: str) -> None:
    _init_sessions()
    _save_active_session()

    with st.sidebar:
        st.markdown(f"### {t('sessions_heading', lang)}")
        st.caption(t("sessions_sub", lang))
        if st.button(t("sessions_new_button", lang), key="saathi_new_session_button", use_container_width=True):
            _create_session()
            st.rerun()

        st.divider()
        active_id = st.session_state[ACTIVE_SESSION_KEY]
        for session_id, meta in st.session_state[SESSIONS_KEY].items():
            title = meta.get("title", session_id)
            label = f"{'● ' if session_id == active_id else ''}{title}"
            if st.button(
                label,
                key=f"saathi_session_button_{session_id}",
                disabled=session_id == active_id,
                use_container_width=True,
            ):
                _switch_session(session_id)
                st.rerun()

        st.caption(t("sessions_zero_persistence", lang))

        # --- "What happens to my data?" privacy explainer ---
        st.divider()
        with st.expander(t("data_privacy_heading", lang)):
            st.markdown(t("data_privacy_body", lang))


def _language_selector() -> str:
    codes = list(LANGUAGES.keys())
    labels = [native_label(c) for c in codes]
    current = st.session_state[LANGUAGE_KEY]
    idx = codes.index(current) if current in codes else 0
    label = st.selectbox(
        t("language_label", current),
        options=labels,
        index=idx,
        key="saathi_language_select",
    )
    new_code = codes[labels.index(label)]
    if new_code != st.session_state[LANGUAGE_KEY]:
        st.session_state[LANGUAGE_KEY] = new_code
        st.rerun()
    return st.session_state[LANGUAGE_KEY]


def _header(lang: str) -> None:
    """Render the gradient wordmark band + language selector."""
    cols = st.columns([5, 1])
    with cols[0]:
        title_en = t("app_title", "en")
        title_local = t("app_title", lang)
        tagline = t("tagline", lang)
        # Always show the Devanagari wordmark alongside the English one so the
        # brand reads cross-lingually — even when the UI is set to Tamil or Telugu,
        # "Saathi / साथी" is instantly recognisable.
        wordmark_suffix = ""
        if title_en != title_local:
            wordmark_suffix = f"<span class='saathi-devanagari'>· {title_local}</span>"
        elif lang == "en":
            wordmark_suffix = "<span class='saathi-devanagari'>· साथी</span>"
        st.markdown(
            f"""
            <div class='saathi-header'>
              <h1 class='saathi-wordmark'>🫂 {title_en} {wordmark_suffix}</h1>
              <p class='saathi-sub'>{tagline}</p>
            </div>
            """,
            unsafe_allow_html=True,
        )
    with cols[1]:
        _language_selector()


def _banners(lang: str) -> None:
    st.info(t("not_therapy_banner", lang))


def _consent_gate(lang: str) -> bool:
    """Display a one-time consent acknowledgment before the main UI.

    Returns True if the user has already consented in this session.
    """
    if st.session_state.get(CONSENT_KEY):
        return True

    st.markdown(_CUSTOM_CSS, unsafe_allow_html=True)
    # Show a minimal header even on the consent screen
    _header(lang)

    st.markdown(
        f"""
        <div class='saathi-consent'>
            <h2>{t('consent_title', lang)}</h2>
            <p style="font-size: 1.05rem; line-height: 1.7; color: #374151;">
                {t('consent_body', lang)}
            </p>
        </div>
        """,
        unsafe_allow_html=True,
    )

    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        if st.button(
            t("consent_button", lang),
            key="saathi_consent_button",
            use_container_width=True,
            type="primary",
        ):
            st.session_state[CONSENT_KEY] = True
            st.rerun()

    return False


def _footer(lang: str) -> None:
    st.divider()
    cols = st.columns([3, 1])
    with cols[0]:
        st.caption(t("footer_disclaimer", lang))
    with cols[1]:
        provider_label = get_active_provider_label()
        st.caption(
            f"_{provider_label} · Built for Anthropic × IIT Delhi 2026 · "
            "[source](https://github.com/samarth2018/Hackathons) · MIT licensed_"
        )


def main() -> None:
    _init_language()
    lang = st.session_state[LANGUAGE_KEY]

    # --- Consent gate (one-time per session) ---
    if not _consent_gate(lang):
        return  # Consent not yet given; stop here

    # --- Main UI (only after consent) ---
    st.markdown(_CUSTOM_CSS, unsafe_allow_html=True)
    _render_session_sidebar(lang)
    _header(lang)
    _banners(lang)

    tab_chat, tab_legal, tab_journal, tab_patterns, tab_soothe = st.tabs(
        [
            t("tab_chat", lang),
            t("tab_legal", lang),
            t("tab_journal", lang),
            t("tab_patterns", lang),
            t("tab_soothe", lang),
        ]
    )

    with tab_chat:
        saathi_chat.render(lang)
    with tab_legal:
        legal_aid.render(lang)
    with tab_journal:
        cognitive_journal.render(lang)
    with tab_patterns:
        my_patterns.render(lang)
    with tab_soothe:
        soothe_poetry.render(lang)

    _footer(lang)


if __name__ == "__main__":
    main()