package server

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"cc.tiredsysadmin.keeporg/internal/secret"
	"cc.tiredsysadmin.keeporg/internal/store"
)

// OIDCConfig configures single sign-on. The server acts as a confidential
// client: it exchanges the authorization code over TLS directly with the
// provider's token endpoint and then reads the identity from the userinfo
// endpoint, so no local JWT verification is needed.
type OIDCConfig struct {
	Issuer       string `json:"issuer"`
	ClientID     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	Scopes       string `json:"scopes"`
	// AutoCreate provisions an account on first successful sign-in.
	AutoCreate bool `json:"auto_create"`
	// DefaultGroupID is the group auto-created users join.
	DefaultGroupID string `json:"default_group_id"`
	// LinkVerifiedEmail lets an identity claim an existing, unlinked local
	// account whose email it shares, when the provider says the email is
	// verified. Off by default: whoever controls an identity provider's
	// email field then controls those accounts. Accounts can always be
	// linked explicitly from their Account page.
	LinkVerifiedEmail bool `json:"link_verified_email"`
}

func (c OIDCConfig) Enabled() bool {
	return c.Issuer != "" && c.ClientID != "" && c.ClientSecret != ""
}

type oidcMetadata struct {
	AuthorizationEndpoint string `json:"authorization_endpoint"`
	TokenEndpoint         string `json:"token_endpoint"`
	UserinfoEndpoint      string `json:"userinfo_endpoint"`
}

type oidcProvider struct {
	cfg         OIDCConfig
	redirectURI string
	http        *http.Client

	mu   sync.Mutex
	meta *oidcMetadata
}

func newOIDCProvider(cfg OIDCConfig, redirectURI string) *oidcProvider {
	if cfg.Scopes == "" {
		cfg.Scopes = "openid profile email"
	}
	return &oidcProvider{cfg: cfg, redirectURI: redirectURI,
		http: &http.Client{Timeout: 15 * time.Second}}
}

// metadata fetches and caches the provider's discovery document.
func (p *oidcProvider) metadata() (*oidcMetadata, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.meta != nil {
		return p.meta, nil
	}
	u := strings.TrimRight(p.cfg.Issuer, "/") + "/.well-known/openid-configuration"
	resp, err := p.http.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("oidc: discovery %s: %s", u, resp.Status)
	}
	var m oidcMetadata
	if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&m); err != nil {
		return nil, err
	}
	if m.AuthorizationEndpoint == "" || m.TokenEndpoint == "" || m.UserinfoEndpoint == "" {
		return nil, errors.New("oidc: discovery document is missing endpoints")
	}
	p.meta = &m
	return p.meta, nil
}

type oidcIdentity struct {
	Subject       string   `json:"sub"`
	Email         string   `json:"email"`
	EmailVerified flexBool `json:"email_verified"`
	Username      string   `json:"preferred_username"`
	Name          string   `json:"name"`
}

// flexBool reads a JSON boolean that some providers send as a string.
type flexBool bool

func (b *flexBool) UnmarshalJSON(data []byte) error {
	v := strings.Trim(string(data), `"`)
	*b = flexBool(v == "true")
	return nil
}

// exchange turns an authorization code into the caller's identity.
func (p *oidcProvider) exchange(code, verifier string) (*oidcIdentity, error) {
	meta, err := p.metadata()
	if err != nil {
		return nil, err
	}
	form := url.Values{
		"grant_type":   {"authorization_code"},
		"code":         {code},
		"redirect_uri": {p.redirectURI},
		// PKCE: only the party that started the flow can spend the code.
		"code_verifier": {verifier},
	}
	req, err := http.NewRequest(http.MethodPost, meta.TokenEndpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth(url.QueryEscape(p.cfg.ClientID), url.QueryEscape(p.cfg.ClientSecret))
	resp, err := p.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<12))
		return nil, fmt.Errorf("oidc: token endpoint: %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	var tok struct {
		AccessToken string `json:"access_token"`
	}
	if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&tok); err != nil {
		return nil, err
	}
	if tok.AccessToken == "" {
		return nil, errors.New("oidc: token response carried no access token")
	}
	return p.userinfo(meta.UserinfoEndpoint, tok.AccessToken)
}

func (p *oidcProvider) userinfo(endpoint, accessToken string) (*oidcIdentity, error) {
	req, err := http.NewRequest(http.MethodGet, endpoint, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	resp, err := p.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("oidc: userinfo: %s", resp.Status)
	}
	var id oidcIdentity
	if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&id); err != nil {
		return nil, err
	}
	if id.Subject == "" {
		return nil, errors.New("oidc: userinfo carried no subject")
	}
	return &id, nil
}

// oidcCookie binds a single sign-on round trip to the browser that started
// it. The state in the URL alone is not enough: a callback link forwarded to
// someone else, or a flow started by an attacker's page in a victim's
// browser, would otherwise finish in the wrong hands. Lax, because the
// provider's redirect back is a cross-site navigation.
const oidcCookie = "keeporg_oidc"

func (s *Server) setOIDCCookie(w http.ResponseWriter, value string, ttl time.Duration) {
	http.SetCookie(w, &http.Cookie{
		Name: oidcCookie, Value: value, Path: "/oidc", HttpOnly: true,
		Secure:   strings.HasPrefix(s.cfg.ExternalURL, "https://"),
		SameSite: http.SameSiteLaxMode, MaxAge: int(ttl.Seconds()),
	})
}

func oidcCookieMatches(r *http.Request, want string) bool {
	c, err := r.Cookie(oidcCookie)
	return err == nil && want != "" && secret.Equal(c.Value, want)
}

// pkceChallenge is the S256 challenge for a verifier.
func pkceChallenge(verifier string) string {
	sum := sha256.Sum256([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(sum[:])
}

// handleOIDCLogin starts the authorization code flow. It has three uses:
// with ?code= it identifies the person approving a waiting client; with
// ?link=1 it links an identity to the account signed in here; without
// either it signs a browser in to the web interface.
func (s *Server) handleOIDCLogin(w http.ResponseWriter, r *http.Request) {
	if s.oidc == nil {
		s.renderDone(w, "Single sign-on unavailable", "This server has no OIDC provider configured.")
		return
	}
	if wait := s.limits.take("start:" + s.clientIP(r)); wait > 0 {
		s.renderDone(w, "Too many sign-in attempts", (&errThrottled{RetryAfter: wait}).Error())
		return
	}
	meta, err := s.oidc.metadata()
	if err != nil {
		s.log.Printf("oidc discovery: %v", err)
		s.renderDone(w, "Single sign-on unavailable", "The identity provider could not be reached.")
		return
	}
	q := r.URL.Query()
	userCode := strings.ToUpper(strings.TrimSpace(q.Get("code")))
	var linkUser string
	if q.Get("link") == "1" {
		// Linking changes how an account signs in, so it needs that
		// account's session and its CSRF token: a page elsewhere must not
		// be able to link the attacker's identity to the victim's account.
		sess := s.webUser(r)
		if sess == nil || !secret.Equal(sess.csrf, q.Get("csrf")) {
			s.renderDone(w, "Single sign-on not linked", "Start again from your Account page.")
			return
		}
		linkUser = sess.user.ID
	}
	state, verifier := secret.Token(), secret.Token()
	if err := s.st.Update(func(st *store.State) error {
		if userCode == "" {
			if !pruneDevices(st) {
				return errBusy
			}
			st.Devices[state] = &store.DeviceAuth{
				DeviceCode: state, Web: linkUser == "", LinkUserID: linkUser, Label: "web",
				OIDCState: state, PKCEVerifier: verifier,
				CreatedAt: time.Now(), ExpiresAt: time.Now().Add(deviceTTL),
			}
			return nil
		}
		d := st.DeviceByUserCode(userCode)
		if d == nil || time.Now().After(d.ExpiresAt) || d.SessionToken != "" {
			return errors.New("unknown or expired device code")
		}
		d.OIDCState, d.PKCEVerifier = state, verifier
		return nil
	}); err != nil {
		s.renderDevice(w, userCode, err.Error())
		return
	}
	s.setOIDCCookie(w, state, deviceTTL)
	v := url.Values{
		"response_type":         {"code"},
		"client_id":             {s.oidc.cfg.ClientID},
		"redirect_uri":          {s.oidc.redirectURI},
		"scope":                 {s.oidc.cfg.Scopes},
		"state":                 {state},
		"code_challenge":        {pkceChallenge(verifier)},
		"code_challenge_method": {"S256"},
	}
	http.Redirect(w, r, meta.AuthorizationEndpoint+"?"+v.Encode(), http.StatusFound)
}

// handleOIDCCallback finishes the flow. A browser sign-in ends signed in; a
// link ends linked; a waiting client is not approved here but shown to the
// person for an explicit decision, because arriving at this URL proves only
// that a browser with a provider session was sent here, not that its owner
// meant to let a device in.
func (s *Server) handleOIDCCallback(w http.ResponseWriter, r *http.Request) {
	if s.oidc == nil {
		s.renderDone(w, "Single sign-on unavailable", "This server has no OIDC provider configured.")
		return
	}
	q := r.URL.Query()
	if q.Get("error") != "" {
		s.renderDone(w, "Sign-in failed", "The identity provider refused the sign-in.")
		return
	}
	state, code := q.Get("state"), q.Get("code")
	if state == "" || code == "" {
		s.renderDone(w, "Sign-in failed", "The provider did not return a code.")
		return
	}
	if !oidcCookieMatches(r, state) {
		s.renderDone(w, "Sign-in failed",
			"This sign-in was not started in this browser. Start again from the same browser.")
		return
	}
	s.setOIDCCookie(w, "", -time.Second)
	var verifier string
	s.st.Read(func(st *store.State) {
		for _, d := range st.Devices {
			if d.OIDCState != "" && secret.Equal(d.OIDCState, state) && time.Now().Before(d.ExpiresAt) {
				verifier = d.PKCEVerifier
			}
		}
	})
	if verifier == "" {
		s.renderDone(w, "Sign-in failed", "This sign-in request has expired. Start again from the client.")
		return
	}
	id, err := s.oidc.exchange(code, verifier)
	if err != nil {
		s.log.Printf("oidc exchange: %v", err)
		s.renderDone(w, "Sign-in failed", "The identity provider rejected the sign-in.")
		return
	}
	var (
		webToken, linked string
		confirm          *webDeviceConfirm
	)
	if err := s.st.Update(func(st *store.State) error {
		var dev *store.DeviceAuth
		for _, d := range st.Devices {
			if d.OIDCState != "" && secret.Equal(d.OIDCState, state) {
				dev = d
			}
		}
		if dev == nil || time.Now().After(dev.ExpiresAt) {
			return errors.New("this sign-in request has expired")
		}
		dev.OIDCState, dev.PKCEVerifier = "", ""
		if dev.LinkUserID != "" {
			delete(st.Devices, dev.DeviceCode)
			return s.linkIdentity(st, dev.LinkUserID, id, &linked)
		}
		u, err := s.linkOIDCUser(st, id)
		if err != nil {
			return err
		}
		if dev.Web {
			// Nothing is waiting to poll for this one; the browser in front
			// of us is the client.
			webToken = s.newBrowserSession(st, u.ID, r).Token
			delete(st.Devices, dev.DeviceCode)
			return nil
		}
		dev.PendingUserID, dev.ApprovalNonce = u.ID, secret.Token()
		confirm = &webDeviceConfirm{Nonce: dev.ApprovalNonce, User: u.Username, Label: dev.Label,
			Code: dev.UserCode, Started: ago(dev.CreatedAt)}
		return nil
	}); err != nil {
		s.renderDone(w, "Sign-in failed", humanError(err))
		return
	}
	switch {
	case webToken != "":
		s.setWebCookie(w, webToken, s.cfg.SessionTTL)
		http.Redirect(w, r, "/web/vaults", http.StatusFound)
	case linked != "":
		http.Redirect(w, r, "/web/account", http.StatusFound)
	default:
		s.setOIDCCookie(w, confirm.Nonce, deviceTTL)
		render(w, deviceConfirmTmpl, map[string]any{"Title": "Approve a device", "D": confirm})
	}
}

// webDeviceConfirm is the question put to someone who signed in through
// their provider to approve a device.
type webDeviceConfirm struct {
	Nonce, User, Label, Code, Started string
}

// handleOIDCApprove is the person's answer. It needs the nonce from the
// confirmation page and the same nonce in this browser's cookie, so it can
// only come from the page the callback showed in this browser.
func (s *Server) handleOIDCApprove(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		fail(w, http.StatusBadRequest, "bad_request", err.Error())
		return
	}
	nonce := r.FormValue("nonce")
	if !oidcCookieMatches(r, nonce) {
		s.renderDone(w, "Not approved", "This approval was not started in this browser.")
		return
	}
	s.setOIDCCookie(w, "", -time.Second)
	approve := r.FormValue("decision") == "approve"
	var returnURL string
	err := s.st.Update(func(st *store.State) error {
		for code, d := range st.Devices {
			if d.ApprovalNonce == "" || !secret.Equal(d.ApprovalNonce, nonce) {
				continue
			}
			if time.Now().After(d.ExpiresAt) || d.SessionToken != "" {
				return errors.New("this sign-in request has expired")
			}
			if !approve {
				delete(st.Devices, code)
				return nil
			}
			u := st.Users[d.PendingUserID]
			if u == nil || u.Disabled {
				return errors.New("this account cannot sign in")
			}
			d.ApprovalNonce = ""
			d.SessionToken = s.newSession(st, u.ID, d.Label).Token
			returnURL = d.ReturnURL
			return nil
		}
		return errors.New("this sign-in request has expired")
	})
	switch {
	case err != nil:
		s.renderDone(w, "Not approved", humanError(err))
	case !approve:
		s.renderDone(w, "Device refused", "The device was not signed in. If you did not start this, someone tried to.")
	default:
		s.renderApproved(w, returnURL)
	}
}

// linkIdentity links an identity to an account its owner is signed in to.
func (s *Server) linkIdentity(st *store.State, userID string, id *oidcIdentity, linked *string) error {
	u := st.Users[userID]
	if u == nil || u.Disabled {
		return errors.New("this account cannot be linked")
	}
	issuer := strings.TrimRight(s.oidc.cfg.Issuer, "/")
	if other := st.UserByOIDC(issuer, id.Subject); other != nil && other.ID != u.ID {
		return errors.New("this identity is already linked to another account")
	}
	u.OIDCIssuer, u.OIDCSubject = issuer, id.Subject
	*linked = u.Username
	return nil
}

// linkOIDCUser finds the account behind an identity, provisioning one when
// the deployment allows it.
func (s *Server) linkOIDCUser(st *store.State, id *oidcIdentity) (*store.User, error) {
	issuer := strings.TrimRight(s.oidc.cfg.Issuer, "/")
	if u := st.UserByOIDC(issuer, id.Subject); u != nil {
		if u.Disabled {
			return nil, errors.New("this account is disabled")
		}
		return u, nil
	}
	// Never by username: anyone who can choose their preferred_username at
	// the provider could otherwise claim "admin". By email only when the
	// deployment allows it and the provider vouches for the address, and
	// never for a server administrator, whose account is linked from its
	// own Account page or not at all.
	if s.oidc.cfg.LinkVerifiedEmail && bool(id.EmailVerified) && id.Email != "" {
		for _, u := range st.Users {
			if u.OIDCSubject != "" || u.ServerAdmin || u.Disabled {
				continue
			}
			if strings.EqualFold(u.Email, id.Email) {
				u.OIDCIssuer, u.OIDCSubject = issuer, id.Subject
				return u, nil
			}
		}
	}
	if !s.oidc.cfg.AutoCreate {
		return nil, errors.New("no account is linked to this identity; sign in with your password and link single sign-on from your Account page")
	}
	name := id.Username
	if name == "" {
		name = id.Email
	}
	if name == "" {
		name = "sso-" + secret.ID()
	}
	if validUsername(name) != nil || st.UserByName(name) != nil {
		name = "sso-" + secret.ID()
	}
	u := &store.User{ID: secret.ID(), Username: name, Email: id.Email,
		OIDCIssuer: issuer, OIDCSubject: id.Subject, CreatedAt: time.Now()}
	st.Users[u.ID] = u
	if g := s.oidc.cfg.DefaultGroupID; g != "" && st.Groups[g] != nil {
		st.Memberships = append(st.Memberships, store.Membership{UserID: u.ID, GroupID: g, Role: store.RoleMember})
	}
	return u, nil
}