// Advantech SSO + Go (standard library only, Go 1.22+)
//
//	SSO_CLIENT_ID=svc_xxxxxxxx SSO_CLIENT_SECRET=xxxxxxxx \
//	APP_URL=https://myapp.advantech.world go run .      (login.html / index.html in ./public)
//
// Register APP_URL/auth/callback as the Redirect URI of the service.
// Sessions live in memory here; use a database or Redis when you run several instances.
package main

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"
)

const SSO = "https://ae-sso.advantech.com"

var (
	clientID     = os.Getenv("SSO_CLIENT_ID")
	clientSecret = os.Getenv("SSO_CLIENT_SECRET")
	appURL       = strings.TrimRight(os.Getenv("APP_URL"), "/")
	redirectURI  = appURL + "/auth/callback"
)

type User map[string]any // sub, email, name, department, job_title, manager, picture ...

type session struct {
	User    User
	Pending map[string]string // state -> page to return to (several tabs may sign in at once)
	Expires time.Time
}

var (
	mu       sync.Mutex
	sessions = map[string]*session{}
)

func randHex() string { b := make([]byte, 16); rand.Read(b); return hex.EncodeToString(b) }

func getSession(w http.ResponseWriter, r *http.Request) *session {
	mu.Lock()
	defer mu.Unlock()
	if c, err := r.Cookie("sid"); err == nil {
		if s := sessions[c.Value]; s != nil && time.Now().Before(s.Expires) {
			return s
		}
	}
	id, s := randHex(), &session{Expires: time.Now().Add(8 * time.Hour)}
	sessions[id] = s
	http.SetCookie(w, &http.Cookie{Name: "sid", Value: id, Path: "/", HttpOnly: true,
		Secure: strings.HasPrefix(appURL, "https://"), SameSite: http.SameSiteLaxMode})
	return s
}

func safeNext(v string) string {
	if strings.HasPrefix(v, "/") && !strings.HasPrefix(v, "//") {
		return v
	}
	return "/"
}

// ---------- 2. is the user signed in? ----------
func currentUser(w http.ResponseWriter, r *http.Request) User { return getSession(w, r).User } // nil = no

func requireLogin(h http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if currentUser(w, r) == nil {
			if strings.HasPrefix(r.URL.Path, "/api/") {
				http.Error(w, `{"error":"not signed in"}`, http.StatusUnauthorized)
				return
			}
			// try single sign-on first; login.html is shown only if that fails
			http.Redirect(w, r, "/auth/login?"+url.Values{"silent": {"1"}, "next": {r.URL.RequestURI()}}.Encode(), http.StatusFound)
			return
		}
		h(w, r)
	}
}

func main() {
	// ---------- 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.
	http.HandleFunc("GET /auth/login", func(w http.ResponseWriter, r *http.Request) {
		s := getSession(w, r)
		state := randHex()
		mu.Lock()
		if s.Pending == nil || len(s.Pending) >= 5 {
			s.Pending = map[string]string{}
		}
		s.Pending[state] = safeNext(r.URL.Query().Get("next"))
		mu.Unlock()
		q := url.Values{"client_id": {clientID}, "redirect_uri": {redirectURI},
			"response_type": {"code"}, "state": {state}}
		if r.URL.Query().Get("silent") == "1" {
			q.Set("prompt", "none")
		}
		http.Redirect(w, r, SSO+"/authorize?"+q.Encode(), http.StatusFound)
	})

	http.HandleFunc("GET /auth/callback", func(w http.ResponseWriter, r *http.Request) {
		s := getSession(w, r)
		state := r.URL.Query().Get("state")
		mu.Lock()
		next, ok := s.Pending[state]
		delete(s.Pending, state)
		mu.Unlock()
		if !ok { // unknown / reused state
			http.Redirect(w, r, "/login.html?error=state", http.StatusFound)
			return
		}
		if e := r.URL.Query().Get("error"); e != "" { // login_required (silent try) / access_denied
			q := url.Values{"next": {next}}
			if e != "login_required" {
				q.Set("error", e)
			}
			http.Redirect(w, r, "/login.html?"+q.Encode(), http.StatusFound)
			return
		}
		res, err := http.PostForm(SSO+"/token", url.Values{
			"grant_type": {"authorization_code"}, "code": {r.URL.Query().Get("code")},
			"redirect_uri": {redirectURI}, "client_id": {clientID}, "client_secret": {clientSecret}})
		if err != nil || res.StatusCode != http.StatusOK {
			http.Redirect(w, r, "/login.html?error=token", http.StatusFound)
			return
		}
		defer res.Body.Close()
		var body struct{ User User }
		if json.NewDecoder(res.Body).Decode(&body) != nil {
			http.Redirect(w, r, "/login.html?error=token", http.StatusFound)
			return
		}
		mu.Lock()
		s.User = body.User
		mu.Unlock()
		http.Redirect(w, r, next, http.StatusFound)
	})

	http.HandleFunc("GET /auth/logout", func(w http.ResponseWriter, r *http.Request) {
		if c, err := r.Cookie("sid"); err == nil {
			mu.Lock()
			delete(sessions, c.Value)
			mu.Unlock()
		}
		http.Redirect(w, r, SSO+"/logout?"+url.Values{"redirect_uri": {appURL + "/"}}.Encode(), http.StatusFound)
	})

	// ---------- 3. user data (used by advantech-sso.js) ----------
	http.HandleFunc("GET /auth/me", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		u := currentUser(w, r)
		if u == nil {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte(`{"error":"not signed in"}`))
			return
		}
		json.NewEncoder(w).Encode(u)
	})

	// ---------- your app ----------
	http.HandleFunc("GET /{$}", requireLogin(func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "public/index.html")
	}))
	http.HandleFunc("GET /api/hello", requireLogin(func(w http.ResponseWriter, r *http.Request) {
		u := currentUser(w, r)
		json.NewEncoder(w).Encode(map[string]any{"message": "Hello " + u["name"].(string), "department": u["department"]})
	}))
	http.Handle("GET /", http.FileServer(http.Dir("public")))

	log.Fatal(http.ListenAndServe(":8080", nil))
}