modules/auth.py

import json
import os
import logging
import secrets
import tempfile
import time
import streamlit as st
import requests
from urllib.parse import urlencode
from modules.config import GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, DOMINIO_PERMITIDO
from modules.access_log import notificar_tentativa_login

_SESSION_FILE = os.path.join(tempfile.gettempdir(), "pm_assistant_sessions.json")
_SESSION_MAX_AGE = 86400  # 24h


def _set_user_session(email, name, photo, access_token=None, refresh_token=None):
    """Set all user-related session state keys after login."""
    st.session_state.usuario_autenticado = True
    st.session_state.usuario_email = email
    st.session_state.usuario_nome = name
    st.session_state.usuario_foto = photo
    if access_token is not None:
        st.session_state.usuario_access_token = access_token
    if refresh_token is not None:
        st.session_state.usuario_refresh_token = refresh_token


def _get_redirect_uri():
    # Support Render, Hugging Face Spaces, or explicit REDIRECT_URI
    uri = os.getenv("RENDER_EXTERNAL_URL") or os.getenv("REDIRECT_URI")
    if not uri:
        space_host = os.getenv("SPACE_HOST")
        if space_host:
            uri = f"https://{space_host}"
    if uri:
        return uri.rstrip("/")
    return "http://localhost:8501"


# --- Secure session persistence (token-based, not email-based) ---

def _save_session(token, email, refresh_token, name="", photo=""):
    """Save session keyed by random token (not email)."""
    try:
        sessions = {}
        if os.path.exists(_SESSION_FILE):
            with open(_SESSION_FILE, "r") as f:
                sessions = json.load(f)
        # Cleanup expired sessions
        now = time.time()
        sessions = {k: v for k, v in sessions.items()
                    if now - v.get("timestamp", 0) < _SESSION_MAX_AGE}
        sessions[token] = {
            "email": email,
            "refresh_token": refresh_token,
            "name": name,
            "photo": photo,
            "timestamp": now,
        }
        with open(_SESSION_FILE, "w") as f:
            json.dump(sessions, f)
    except Exception as e:
        logging.warning(f"[Auth] Failed to save session: {e}")


def _load_session(token):
    """Load session by its unique token. Returns dict or None."""
    try:
        if not token or not os.path.exists(_SESSION_FILE):
            return None
        with open(_SESSION_FILE, "r") as f:
            sessions = json.load(f)
        data = sessions.get(token)
        if not data:
            return None
        if time.time() - data.get("timestamp", 0) > _SESSION_MAX_AGE:
            return None
        return data
    except Exception:
        return None


def _clear_session(token):
    """Remove a specific session by token."""
    try:
        if not os.path.exists(_SESSION_FILE):
            return
        with open(_SESSION_FILE, "r") as f:
            sessions = json.load(f)
        sessions.pop(token, None)
        with open(_SESSION_FILE, "w") as f:
            json.dump(sessions, f)
    except Exception:
        pass


def _try_auto_login():
    """Try auto-login using session token from query params."""
    token = st.query_params.get("sid")
    if not token:
        return False
    saved = _load_session(token)
    if not saved or not saved.get("refresh_token"):
        return False
    try:
        refresh_resp = requests.post("https://oauth2.googleapis.com/token", data={
            "client_id": GOOGLE_CLIENT_ID,
            "client_secret": GOOGLE_CLIENT_SECRET,
            "refresh_token": saved["refresh_token"],
            "grant_type": "refresh_token",
        }, timeout=10)
        if refresh_resp.status_code != 200:
            return False
        token_data = refresh_resp.json()
        _set_user_session(
            email=saved["email"],
            name=saved.get("name", saved["email"]),
            photo=saved.get("photo", ""),
            access_token=token_data.get("access_token"),
            refresh_token=saved["refresh_token"],
        )
        logging.info(f"[Auth] Auto-login OK for {saved['email']}")
        notificar_tentativa_login(saved["email"], saved.get("name", saved["email"]), "auto_login")
        return True
    except Exception as e:
        logging.warning(f"[Auth] Auto-login failed: {e}")
        return False


def logout():
    """Clear session from server and query params."""
    token = st.query_params.get("sid")
    if token:
        _clear_session(token)
    st.query_params.clear()


def _google_oauth_login():
    """Exibe tela de login e redireciona para Google OAuth."""
    st.markdown("""
    <style>
    section[data-testid="stSidebar"] {display: none !important;}
    header[data-testid="stHeader"] {display: none !important;}
    .block-container {padding-top: 0 !important; max-width: 100% !important;}
    .login-bg {
        min-height: 100vh; display: flex; align-items: center; justify-content: center;
        flex-direction: column; text-align: center; padding: 40px 20px;
        position: fixed; top: 0; left: 0; right: 0; bottom: 0;
        background: white; z-index: 9999;
    }
    .login-logo {
        font-size: 3em; font-weight: 800; color: #FE3E6D; margin-bottom: 8px;
        letter-spacing: -1px;
    }
    .login-subtitle {
        font-size: 1.15em; color: #666; margin-bottom: 48px; line-height: 1.6;
    }
    .google-btn {
        display: inline-flex; align-items: center; gap: 12px;
        background: #FE3E6D; color: white !important; border: none;
        border-radius: 12px; padding: 14px 36px; font-size: 1.05em;
        text-decoration: none !important; cursor: pointer;
        font-weight: 600; transition: all 0.2s;
        box-shadow: 0 4px 14px rgba(254, 62, 109, 0.3);
    }
    .google-btn:hover {
        background: #FF6B8A; transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(254, 62, 109, 0.4);
    }
    .login-footer {font-size: 0.85em; color: #999; margin-top: 32px;}
    </style>
    """, unsafe_allow_html=True)

    params = st.query_params
    code = params.get("code")

    if code:
        token_resp = requests.post("https://oauth2.googleapis.com/token", data={
            "code": code,
            "client_id": GOOGLE_CLIENT_ID,
            "client_secret": GOOGLE_CLIENT_SECRET,
            "redirect_uri": _get_redirect_uri(),
            "grant_type": "authorization_code",
        }, timeout=10)
        if token_resp.status_code == 200:
            token_data = token_resp.json()
            user_resp = requests.get(
                "https://www.googleapis.com/oauth2/v2/userinfo",
                headers={"Authorization": f"Bearer {token_data['access_token']}"},
                timeout=10,
            )
            if user_resp.status_code == 200:
                user_info = user_resp.json()
                email = user_info.get("email", "")
                if email.endswith(f"@{DOMINIO_PERMITIDO}"):
                    _set_user_session(
                        email=email,
                        name=user_info.get("name", email),
                        photo=user_info.get("picture", ""),
                        access_token=token_data.get("access_token"),
                        refresh_token=token_data.get("refresh_token"),
                    )
                    logging.info(f"[Auth] Login OK for {email}")
                    notificar_tentativa_login(email, user_info.get("name", email), "sucesso")
                    # Save session with random token for persistence
                    rt = token_data.get("refresh_token")
                    if rt:
                        sid = secrets.token_urlsafe(32)
                        _save_session(sid, email, rt,
                                      user_info.get("name", email),
                                      user_info.get("picture", ""))
                        st.query_params.clear()
                        st.query_params["sid"] = sid
                    else:
                        st.query_params.clear()
                    st.rerun()
                else:
                    notificar_tentativa_login(email, user_info.get("name", email), "dominio_negado")
                    st.error(f"Acesso negado. Apenas emails @{DOMINIO_PERMITIDO} podem acessar. Voce tentou com: {email}")
                    st.stop()
        else:
            notificar_tentativa_login("desconhecido", "", "erro_oauth")
            st.error("Erro ao autenticar com Google. Tente novamente.")
            st.stop()

    # Try auto-login from session token in query params
    if _try_auto_login():
        st.rerun()

    auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode({
        "client_id": GOOGLE_CLIENT_ID,
        "redirect_uri": _get_redirect_uri(),
        "response_type": "code",
        "scope": "openid email profile https://www.googleapis.com/auth/drive",
        "access_type": "offline",
        "hd": DOMINIO_PERMITIDO,
        "prompt": "consent",
    })

    st.markdown(f"""
    <div class="login-bg">
        <div class="login-logo">cora</div>
        <div style="font-size: 1.4em; color: #2D2D2D; font-weight: 600; margin-bottom: 8px;">PM Assistant</div>
        <div class="login-subtitle">Assistente de produto inteligente<br>Acesso restrito a colaboradores @cora.com.br</div>
        <a href="{auth_url}" class="google-btn">
            <img src="https://www.gstatic.com/firebasejs/ui/2.0.0/images/auth/google.svg" width="22">
            Entrar com Google
        </a>
        <div class="login-footer">Apenas emails @cora.com.br sao aceitos</div>
    </div>
    """, unsafe_allow_html=True)
    st.stop()


def verificar_autenticacao():
    if not GOOGLE_CLIENT_ID or not GOOGLE_CLIENT_SECRET:
        _set_user_session(
            email="dev@local",
            name="Dev Local",
            photo="",
        )
        return
    if not st.session_state.get("usuario_autenticado"):
        _google_oauth_login()