src/accounts/ui.py

"""Streamlit sign-in / sign-up gate and session-cookie glue.

Kept out of the main app file so the auth flow is self-contained. The cookie
component (``extra_streamlit_components``) is imported lazily; if it is not
installed the app still works, but "stay logged in" downgrades to per-session
(a browser refresh will require signing in again).
"""

from __future__ import annotations

import datetime as _dt
import logging
from typing import Any, Dict, Optional

import streamlit as st

from .auth import AuthError

logger = logging.getLogger("agadvisor.accounts.ui")

COOKIE_NAME = "agadvisor_session"
_COOKIE_TTL_DAYS = 30


def _get_cookie_manager():
    """Return a cached CookieManager, or None if the component is unavailable."""
    if "cookie_manager" in st.session_state:
        return st.session_state.cookie_manager
    try:
        import extra_streamlit_components as stx  # lazy

        mgr = stx.CookieManager(key="agadvisor_cookie_mgr")
        st.session_state.cookie_manager = mgr
        return mgr
    except Exception as e:  # not installed / component error
        logger.info("Cookie manager unavailable (%s); using session-only auth.", e)
        st.session_state.cookie_manager = None
        return None


def _set_session_cookie(token: str) -> None:
    mgr = _get_cookie_manager()
    if mgr is None:
        return
    try:
        expires = _dt.datetime.now() + _dt.timedelta(days=_COOKIE_TTL_DAYS)
        mgr.set(COOKIE_NAME, token, expires_at=expires, key="set_session_cookie")
    except Exception as e:
        logger.info("Could not set session cookie: %s", e)


def _clear_session_cookie() -> None:
    mgr = _get_cookie_manager()
    if mgr is None:
        return
    try:
        mgr.delete(COOKIE_NAME, key="del_session_cookie")
    except Exception as e:
        logger.info("Could not clear session cookie: %s", e)


def _read_cookie_token() -> Optional[str]:
    mgr = _get_cookie_manager()
    if mgr is None:
        return None
    try:
        return mgr.get(COOKIE_NAME)
    except Exception:
        return None


def _resolve_user(service) -> Optional[Dict[str, Any]]:
    """Return the logged-in user from session_state, else from a valid cookie."""
    if st.session_state.get("user"):
        return st.session_state["user"]
    token = _read_cookie_token()
    if token:
        user = service.user_from_token(token)
        if user:
            st.session_state["user"] = {"id": user["id"], "username": user["username"]}
            return st.session_state["user"]
    return None


def logout(service=None) -> None:
    """Flush any pending sync, clear the cookie and session, and rerun."""
    try:
        if service is not None and getattr(service, "sync", None) is not None:
            service.sync.flush()
    except Exception:
        pass
    _clear_session_cookie()
    for k in ("user", "chats", "current_chat_id", "chat_counter"):
        st.session_state.pop(k, None)
    st.rerun()


def render_login_gate(service) -> Optional[Dict[str, Any]]:
    """Render the sign-in / sign-up UI. Returns the user on success, else None
    (call ``st.stop()`` in the caller when None)."""
    st.markdown("### 🌿 Welcome to AgAdvisor")
    st.caption("Sign in to save your conversations. Your chat history is private to your account.")

    tab_login, tab_signup = st.tabs(["Sign in", "Create account"])

    with tab_login:
        with st.form("login_form", clear_on_submit=False):
            username = st.text_input("Username", key="login_username")
            password = st.text_input("Password", type="password", key="login_password")
            submitted = st.form_submit_button("Sign in", use_container_width=True, type="primary")
        if submitted:
            try:
                user = service.login(username, password)
            except AuthError as e:
                st.error(str(e))
            else:
                st.session_state["user"] = user
                _set_session_cookie(service.issue_token(user["id"]))
                st.success(f"Welcome back, {user['username']}!")
                st.rerun()

    with tab_signup:
        with st.form("signup_form", clear_on_submit=False):
            new_username = st.text_input("Choose a username", key="signup_username")
            new_password = st.text_input("Choose a password", type="password", key="signup_password")
            confirm = st.text_input("Confirm password", type="password", key="signup_confirm")
            submitted = st.form_submit_button("Create account", use_container_width=True, type="primary")
        if submitted:
            if new_password != confirm:
                st.error("Passwords do not match.")
            else:
                try:
                    user = service.signup(new_username, new_password)
                except AuthError as e:
                    st.error(str(e))
                else:
                    st.session_state["user"] = user
                    _set_session_cookie(service.issue_token(user["id"]))
                    st.success(f"Account created — welcome, {user['username']}!")
                    st.rerun()

    return None


def require_auth(service) -> Dict[str, Any]:
    """Gate the app: return the current user, or render the login gate and stop."""
    user = _resolve_user(service)
    if user:
        return user
    render_login_gate(service)
    st.stop()