// Copyright (C) NHR@FAU, University Erlangen-Nuremberg.
// All rights reserved. This file is part of cc-backend.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package auth

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strings"
	"time"

	"github.com/ClusterCockpit/cc-backend/internal/repository"
	cclog "github.com/ClusterCockpit/cc-lib/v2/ccLogger"
	"github.com/ClusterCockpit/cc-lib/v2/schema"
	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/go-chi/chi/v5"
	"golang.org/x/oauth2"
)

type OpenIDConfig struct {
	Provider          string `json:"provider"`
	SyncUserOnLogin   bool   `json:"sync-user-on-login"`
	UpdateUserOnLogin bool   `json:"update-user-on-login"`

	// OAuth2 client ID for the OIDC provider.
	// Overridden by the OID_CLIENT_ID environment variable when set.
	ClientID string `json:"client-id"`

	// OAuth2 client secret for the OIDC provider.
	// Overridden by the OID_CLIENT_SECRET environment variable when set.
	ClientSecret string `json:"client-secret"`

	// Maps an OIDC role/group claim value to a CC role (admin/support/api/manager/user).
	// This is the sole source of roles: a token role grants a CC role only if it is
	// listed here. Unmapped token roles are ignored (no identity fallback), so literal
	// CC role names must be mapped explicitly. Users without any mapped role get "user".
	RoleMapping map[string]string `json:"role-mapping"`
}

type OIDC struct {
	client         *oauth2.Config
	provider       *oidc.Provider
	authentication *Authentication
	clientID       string
	// roleMapping is the validated subset of OpenIDConfig.RoleMapping
	// (IdP role/group name -> CC role).
	roleMapping map[string]string
}

func randString(nByte int) (string, error) {
	b := make([]byte, nByte)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

func setCallbackCookie(w http.ResponseWriter, r *http.Request, name, value string) {
	c := &http.Cookie{
		Name:     name,
		Value:    value,
		MaxAge:   int(time.Hour.Seconds()),
		Secure:   r.TLS != nil,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	}
	http.SetCookie(w, c)
}

// NewOIDC creates a new OIDC authenticator with the configured provider
func NewOIDC(a *Authentication) *OIDC {
	// Use context with timeout for provider initialization
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	provider, err := oidc.NewProvider(ctx, Keys.OpenIDConfig.Provider)
	if err != nil {
		cclog.Fatal(err)
	}
	clientID := secretFromEnv("OID_CLIENT_ID", Keys.OpenIDConfig.ClientID)
	if clientID == "" {
		cclog.Warn("OIDC client ID not configured ('client-id' in config or 'OID_CLIENT_ID' env): Open ID connect auth will not work")
	}
	clientSecret := secretFromEnv("OID_CLIENT_SECRET", Keys.OpenIDConfig.ClientSecret)
	if clientSecret == "" {
		cclog.Warn("OIDC client secret not configured ('client-secret' in config or 'OID_CLIENT_SECRET' env): Open ID connect auth will not work")
	}

	client := &oauth2.Config{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		Endpoint:     provider.Endpoint(),
		Scopes:       []string{oidc.ScopeOpenID, "profile", "roles"},
	}

	// Validate the optional role mapping once at startup. Invalid targets are
	// dropped with a warning rather than failing startup. IdP names (keys) are
	// kept verbatim so they match the raw token claim values.
	roleMapping := make(map[string]string)
	for name, ccRole := range Keys.OpenIDConfig.RoleMapping {
		role := strings.ToLower(ccRole)
		if !schema.IsValidRole(role) || role == schema.GetRoleString(schema.RoleAnonymous) {
			cclog.Warnf("OIDC: ignoring role-mapping '%s' -> '%s': invalid or non-assignable target role", name, ccRole)
			continue
		}
		roleMapping[name] = role
	}

	oa := &OIDC{provider: provider, client: client, clientID: clientID, authentication: a, roleMapping: roleMapping}

	return oa
}

// mapOIDCRoles translates raw OIDC role/group names into the CC role set using
// only the configured mapping (IdP name -> CC role). Unmapped names are ignored.
// Always returns at least [user]. The result is deduplicated and sorted.
func mapOIDCRoles(oidcRoles []string, mapping map[string]string) []string {
	roleSet := make(map[string]bool)
	for _, r := range oidcRoles {
		if cc, ok := mapping[r]; ok {
			roleSet[cc] = true
		}
	}

	if len(roleSet) == 0 {
		return []string{schema.GetRoleString(schema.RoleUser)}
	}

	roles := make([]string, 0, len(roleSet))
	for role := range roleSet {
		roles = append(roles, role)
	}
	sort.Strings(roles)
	return roles
}

func (oa *OIDC) RegisterEndpoints(r chi.Router) {
	r.HandleFunc("/oidc-login", oa.OAuth2Login)
	r.HandleFunc("/oidc-callback", oa.OAuth2Callback)
}

func (oa *OIDC) OAuth2Callback(rw http.ResponseWriter, r *http.Request) {
	c, err := r.Cookie("state")
	if err != nil {
		http.Error(rw, "state cookie not found", http.StatusBadRequest)
		return
	}
	state := c.Value

	c, err = r.Cookie("verifier")
	if err != nil {
		http.Error(rw, "verifier cookie not found", http.StatusBadRequest)
		return
	}
	codeVerifier := c.Value

	_ = r.ParseForm()
	if r.Form.Get("state") != state {
		http.Error(rw, "State invalid", http.StatusBadRequest)
		return
	}
	code := r.Form.Get("code")
	if code == "" {
		http.Error(rw, "Code not found", http.StatusBadRequest)
		return
	}
	// Exchange authorization code for token with timeout
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	token, err := oa.client.Exchange(ctx, code, oauth2.VerifierOption(codeVerifier))
	if err != nil {
		cclog.Errorf("token exchange failed: %s", err.Error())
		http.Error(rw, "Authentication failed during token exchange", http.StatusInternalServerError)
		return
	}

	// Get user info from OIDC provider with same timeout
	userInfo, err := oa.provider.UserInfo(ctx, oauth2.StaticTokenSource(token))
	if err != nil {
		cclog.Errorf("failed to get userinfo: %s", err.Error())
		http.Error(rw, "Failed to retrieve user information", http.StatusInternalServerError)
		return
	}

	// Verify ID token and nonce to prevent replay attacks
	rawIDToken, ok := token.Extra("id_token").(string)
	if !ok {
		http.Error(rw, "ID token not found in response", http.StatusInternalServerError)
		return
	}

	nonceCookie, err := r.Cookie("nonce")
	if err != nil {
		http.Error(rw, "nonce cookie not found", http.StatusBadRequest)
		return
	}

	verifier := oa.provider.Verifier(&oidc.Config{ClientID: oa.clientID})
	idToken, err := verifier.Verify(ctx, rawIDToken)
	if err != nil {
		cclog.Errorf("ID token verification failed: %s", err.Error())
		http.Error(rw, "ID token verification failed", http.StatusInternalServerError)
		return
	}

	if idToken.Nonce != nonceCookie.Value {
		http.Error(rw, "Nonce mismatch", http.StatusBadRequest)
		return
	}

	// projects is populated below from ID token claims
	var projects []string

	// Extract profile claims from userinfo (username, name)
	var userInfoClaims struct {
		Username string `json:"preferred_username"`
		Name     string `json:"name"`
	}
	if err := userInfo.Claims(&userInfoClaims); err != nil {
		cclog.Errorf("failed to extract userinfo claims: %s", err.Error())
		http.Error(rw, "Failed to extract user claims", http.StatusInternalServerError)
		return
	}

	// Extract role claims from the ID token.
	// Keycloak includes realm_access and resource_access in the ID token (JWT),
	// but NOT in the UserInfo endpoint response by default.
	var idTokenClaims struct {
		Username string `json:"preferred_username"`
		Name     string `json:"name"`
		// Keycloak realm-level roles
		RealmAccess struct {
			Roles []string `json:"roles"`
		} `json:"realm_access"`
		// Keycloak client-level roles: map from client-id to role list
		ResourceAccess map[string]struct {
			Roles []string `json:"roles"`
		} `json:"resource_access"`
		// Custom multi-valued user attribute mapped via a Keycloak User Attribute mapper
		Projects []string `json:"projects"`
	}
	if err := idToken.Claims(&idTokenClaims); err != nil {
		cclog.Errorf("failed to extract ID token claims: %s", err.Error())
		http.Error(rw, "Failed to extract ID token claims", http.StatusInternalServerError)
		return
	}

	cclog.Debugf("OIDC userinfo claims: username=%q name=%q", userInfoClaims.Username, userInfoClaims.Name)
	cclog.Debugf("OIDC ID token realm_access roles: %v", idTokenClaims.RealmAccess.Roles)
	cclog.Debugf("OIDC ID token resource_access: %v", idTokenClaims.ResourceAccess)
	cclog.Debugf("OIDC ID token projects: %v", idTokenClaims.Projects)

	projects = idTokenClaims.Projects
	if projects == nil {
		projects = []string{}
	}

	// Prefer username from userInfo; fall back to ID token claim
	username := userInfoClaims.Username
	if username == "" {
		username = idTokenClaims.Username
	}
	name := userInfoClaims.Name
	if name == "" {
		name = idTokenClaims.Name
	}

	if username == "" {
		http.Error(rw, "Username claim missing from OIDC provider", http.StatusBadRequest)
		return
	}

	// Collect roles from realm_access (realm roles) in the ID token
	oidcRoles := append([]string{}, idTokenClaims.RealmAccess.Roles...)

	// Also collect roles from resource_access (client roles) for all clients
	for clientID, access := range idTokenClaims.ResourceAccess {
		cclog.Debugf("OIDC ID token resource_access[%q] roles: %v", clientID, access.Roles)
		oidcRoles = append(oidcRoles, access.Roles...)
	}

	roles := mapOIDCRoles(oidcRoles, oa.roleMapping)

	user := &schema.User{
		Username:   username,
		Name:       name,
		Roles:      roles,
		Projects:   projects,
		AuthSource: schema.AuthViaOIDC,
	}

	if Keys.OpenIDConfig.SyncUserOnLogin || Keys.OpenIDConfig.UpdateUserOnLogin {
		handleOIDCUser(user)
	}

	if err := oa.authentication.SaveSession(rw, r, user); err != nil {
		cclog.Errorf("session save failed for user %q: %s", user.Username, err.Error())
		http.Error(rw, "Failed to create session", http.StatusInternalServerError)
		return
	}
	cclog.Infof("login successful: user: %#v (roles: %v, projects: %v)", user.Username, user.Roles, user.Projects)
	userCtx := context.WithValue(r.Context(), repository.ContextUserKey, user)
	http.RedirectHandler("/", http.StatusTemporaryRedirect).ServeHTTP(rw, r.WithContext(userCtx))
}

func (oa *OIDC) OAuth2Login(rw http.ResponseWriter, r *http.Request) {
	state, err := randString(16)
	if err != nil {
		http.Error(rw, "Internal error", http.StatusInternalServerError)
		return
	}
	setCallbackCookie(rw, r, "state", state)

	// use PKCE to protect against CSRF attacks
	codeVerifier := oauth2.GenerateVerifier()
	setCallbackCookie(rw, r, "verifier", codeVerifier)

	// Generate nonce for ID token replay protection
	nonce, err := randString(16)
	if err != nil {
		http.Error(rw, "Internal error", http.StatusInternalServerError)
		return
	}
	setCallbackCookie(rw, r, "nonce", nonce)

	// Build redirect URL from the incoming request
	scheme := "https"
	if r.TLS == nil && r.Header.Get("X-Forwarded-Proto") != "https" {
		scheme = "http"
	}
	oa.client.RedirectURL = fmt.Sprintf("%s://%s/oidc-callback", scheme, r.Host)

	// Redirect user to consent page to ask for permission
	url := oa.client.AuthCodeURL(state, oauth2.AccessTypeOffline,
		oauth2.S256ChallengeOption(codeVerifier),
		oidc.Nonce(nonce))
	http.Redirect(rw, r, url, http.StatusFound)
}
