# Advantech SSO + Flask  —  pip install flask requests
#
#   export SSO_CLIENT_ID=svc_xxxxxxxx
#   export SSO_CLIENT_SECRET=xxxxxxxx
#   export APP_URL=https://myapp.advantech.world      # this site; register APP_URL/auth/callback as Redirect URI
#   export FLASK_SECRET_KEY=$(python -c "import secrets;print(secrets.token_hex(32))")
#   python flask_app.py           (put login.html / index.html in ./public)
import os, secrets
from functools import wraps
from urllib.parse import urlencode

import requests
from flask import Flask, abort, jsonify, redirect, request, send_from_directory, session

SSO = "https://ae-sso.advantech.com"
CLIENT_ID = os.environ["SSO_CLIENT_ID"]
CLIENT_SECRET = os.environ["SSO_CLIENT_SECRET"]
APP_URL = os.environ["APP_URL"].rstrip("/")
REDIRECT_URI = APP_URL + "/auth/callback"

# the session is a signed cookie (max ~4 KB), so keep only the fields you use
KEEP = ("sub", "email", "name", "given_name", "family_name", "job_title", "department",
        "company", "office_location", "employee_id", "manager", "picture")

app = Flask(__name__, static_folder="public", static_url_path="")
app.secret_key = os.environ["FLASK_SECRET_KEY"]
app.config.update(SESSION_COOKIE_SECURE=APP_URL.startswith("https://"),
                  SESSION_COOKIE_HTTPONLY=True, SESSION_COOKIE_SAMESITE="Lax")


def safe_next(v):
    return v if v and v.startswith("/") and not v.startswith("//") else "/"


# ---------- 1. sign in ----------
# /auth/login?silent=1 adds prompt=none: if the user already signed in to ANY
# service, the SSO returns straight away (no page, no click); otherwise it comes
# back with error=login_required and we show login.html.
@app.get("/auth/login")
def login():
    state = secrets.token_urlsafe(16)
    # remember several pending sign-ins (two tabs, a double click ...), newest 5
    pending = dict(list(session.get("pending", {}).items())[-4:])
    pending[state] = safe_next(request.args.get("next"))
    session["pending"] = pending
    q = {"client_id": CLIENT_ID, "redirect_uri": REDIRECT_URI,
         "response_type": "code", "state": state}
    if request.args.get("silent") == "1":
        q["prompt"] = "none"
    return redirect(f"{SSO}/authorize?{urlencode(q)}")


@app.get("/auth/callback")
def callback():
    pending = session.get("pending", {})
    next_url = pending.pop(request.args.get("state", ""), None)
    session["pending"] = pending
    if next_url is None:                 # unknown / reused state
        return redirect("/login.html?error=state")
    if request.args.get("error"):      # login_required (silent try) / access_denied
        q = {"next": next_url}
        if request.args["error"] != "login_required":
            q["error"] = request.args["error"]
        return redirect("/login.html?" + urlencode(q))
    r = requests.post(f"{SSO}/token", data={
        "grant_type": "authorization_code",
        "code": request.args.get("code", ""),
        "redirect_uri": REDIRECT_URI,
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
    }, timeout=10)
    if r.status_code != 200:
        return redirect("/login.html?error=token")
    user = r.json()["user"]
    session["user"] = {k: user[k] for k in KEEP if k in user}
    return redirect(next_url)


@app.get("/auth/logout")
def logout():
    session.clear()
    return redirect(f"{SSO}/logout?" + urlencode({"redirect_uri": APP_URL + "/"}))


# ---------- 2. is the user signed in? ----------
def current_user():
    return session.get("user")          # None = not signed in


def login_required(view):
    @wraps(view)
    def wrapper(*args, **kwargs):
        if current_user() is None:
            if request.path.startswith("/api/"):
                abort(401)
            # try single sign-on first; login.html is shown only if that fails
            return redirect("/auth/login?" + urlencode({"silent": 1, "next": request.full_path.rstrip("?")}))
        return view(*args, **kwargs)
    return wrapper


# ---------- 3. user data (used by advantech-sso.js) ----------
@app.get("/auth/me")
def me():
    u = current_user()
    return (jsonify(u), 200) if u else (jsonify(error="not signed in"), 401)


# ---------- your app ----------
@app.get("/")
@login_required
def index():
    return send_from_directory("public", "index.html")


@app.get("/api/hello")
@login_required
def hello():
    u = current_user()
    return {"message": f"Hello {u['name']}", "department": u.get("department")}


if __name__ == "__main__":
    app.run(port=5000)