from urllib.parse import urlparse

from flask import flash, redirect, session, url_for
from flask_jwt_extended import create_access_token


def _safe_next(relay_state: str | None) -> str | None:
    if not relay_state:
        return None
    parsed = urlparse(relay_state)
    if parsed.netloc:
        return None
    if not relay_state.startswith('/'):
        return None
    return relay_state


def establish_user_session(user, relay_state: str | None = None, auth_method: str = 'password'):
    """Create a browser session + JWT for the provided user."""
    session['user_id'] = user.id
    session['username'] = user.username
    session['is_admin'] = user.is_admin
    session['auth_method'] = auth_method

    access_token = create_access_token(identity=str(user.id))
    session['access_token'] = access_token

    if user.password_change_required and auth_method != 'saml':
        flash('You must change your password before continuing', 'warning')
        target = url_for('web.change_password')
    else:
        target = _safe_next(relay_state) or url_for('web.dashboard')
        flash('Logged in successfully', 'success')

    response = redirect(target)
    response.set_cookie('access_token', access_token, httponly=False, secure=False, samesite='Lax')
    return response