package idp

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/server"
	"github.com/dexidp/dex/server/connectors"
	dexstorage "github.com/dexidp/dex/storage"

	"tug/internal/service/identity"
	"tug/internal/service/refusalcap"
	"tug/internal/storage"
)

const (
	ConnectorID   = storage.DexConnectorID
	ConnectorType = "unified"
	ValueParam    = "tug"
	valueLength   = 32
	userIDLength  = 36
	standInDomain = "tug.invalid"
)

var (
	ErrWrongConnector     = errors.New("idp: auth request belongs to another connector")
	ErrAuthRequestExpired = errors.New("idp: auth request expired")
	ErrAlreadyUsed        = errors.New("idp: login value already used")
	ErrNoPending          = errors.New("idp: no pending login on this auth request")
	ErrValueExpired       = errors.New("idp: login value expired")
	ErrValueMismatch      = errors.New("idp: login value does not match")
	ErrValueMalformed     = errors.New("idp: login value missing or malformed")
	ErrUnregistered       = errors.New("idp: no dependencies registered for this connector key")
	ErrGrantMalformed     = errors.New("idp: refresh grant carries no credentials epoch")
	ErrGrantRevoked       = errors.New("idp: credentials revoked since this authentication")
)

type deps struct {
	identity      *identity.Service
	store         *storage.Store
	dex           dexstorage.Storage
	now           func() time.Time
	issuer        string
	accountOrigin string
	accountDomain string
	refusals      *refusalcap.Cap
}

var (
	registry     sync.Map
	registerOnce sync.Once
)

func registerConnectorType() {
	registerOnce.Do(func() {
		server.ConnectorsConfig[ConnectorType] = func() connectors.ConnectorConfig { return new(connectorConfig) }
	})
}

type connectorConfig struct {
	Key string `json:"key"`
}

func (c *connectorConfig) Open(id string, logger *slog.Logger) (connector.Connector, error) {
	v, ok := registry.Load(c.Key)
	if !ok {
		return nil, ErrUnregistered
	}
	return &unified{deps: v.(*deps)}, nil
}

type unified struct {
	deps *deps
}

var (
	_ connector.CallbackConnector = (*unified)(nil)
	_ connector.RefreshConnector  = (*unified)(nil)
)

type pending struct {
	Binding        string `json:"k,omitempty"`
	Hint           string `json:"h,omitempty"`
	UserID         string `json:"u,omitempty"`
	Digest         []byte `json:"d,omitempty"`
	Expiry         int64  `json:"x,omitempty"`
	Epoch          *int64 `json:"e,omitempty"`
	Browser        string `json:"b,omitempty"`
	RequiredDigest []byte `json:"r,omitempty"`
	RequiredExpiry int64  `json:"rx,omitempty"`
}

func encodePending(p pending) []byte {
	b, _ := json.Marshal(p)
	return b
}

func parsePending(data []byte) (pending, error) {
	data = bytes.TrimSpace(data)
	if len(data) >= userIDLength && data[0] != '{' {
		return pending{}, ErrAlreadyUsed
	}
	p, err := record(data)
	if err != nil {
		return pending{}, ErrNoPending
	}
	if len(p.UserID) != userIDLength || len(p.Digest) != sha256.Size || p.Expiry <= 0 || p.Epoch == nil || *p.Epoch < 0 || (p.Browser != "" && !storage.ValidBrowser(p.Browser)) {
		return pending{}, ErrNoPending
	}
	return p, nil
}

func (u *unified) LoginURL(s connector.Scopes, callbackURL, state string) (string, []byte, error) {
	q := url.Values{}
	q.Set("state", state)
	return u.deps.accountOrigin + "/login?" + q.Encode(), nil, nil
}

func (u *unified) check(a dexstorage.AuthRequest, value []byte) (pending, error) {
	now := u.deps.now()
	if a.ConnectorID != ConnectorID {
		return pending{}, ErrWrongConnector
	}
	if now.After(a.Expiry) {
		return pending{}, ErrAuthRequestExpired
	}
	p, err := parsePending(a.ConnectorData)
	if err != nil {
		return pending{}, err
	}
	if err := checkValue(value, p.Digest, p.Expiry, now); err != nil {
		return pending{}, err
	}
	return p, nil
}

func (u *unified) HandleCallback(s connector.Scopes, connData []byte, r *http.Request) (connector.Identity, error) {
	ctx := r.Context()
	q, err := url.ParseQuery(r.URL.RawQuery)
	if err != nil {
		return connector.Identity{}, fmt.Errorf("%w: %w", ErrValueMalformed, err)
	}
	value, err := base64.RawURLEncoding.Strict().DecodeString(q.Get(ValueParam))
	if err != nil || len(value) != valueLength {
		return connector.Identity{}, ErrValueMalformed
	}
	var userID, browser, clientID string
	var epoch int64
	err = u.deps.dex.UpdateAuthRequest(ctx, q.Get("state"), func(a dexstorage.AuthRequest) (dexstorage.AuthRequest, error) {
		p, err := u.check(a, value)
		if err != nil {
			return a, err
		}
		userID, epoch, browser, clientID = p.UserID, *p.Epoch, p.Browser, a.ClientID
		a.ConnectorData = encodeGrant(p.UserID, *p.Epoch, p.Browser)
		return a, nil
	})
	if err != nil {
		return connector.Identity{}, fmt.Errorf("idp: spend login value: %w", err)
	}
	ident, err := u.identity(ctx, s, userID)
	if err != nil {
		return connector.Identity{}, err
	}
	ident.ConnectorData = encodeGrant(userID, epoch, browser)
	if err := u.deps.store.RecordLoginAudited(ctx, userID, u.deps.now(), loginSuccess(userID, clientID)); err != nil {
		return connector.Identity{}, fmt.Errorf("idp: record login: %w", err)
	}
	return ident, nil
}

func (u *unified) identity(ctx context.Context, s connector.Scopes, userID string) (connector.Identity, error) {
	claims, err := u.deps.identity.Claims(ctx, userID, s.Groups)
	if err != nil {
		return connector.Identity{}, fmt.Errorf("idp: load user: %w", err)
	}
	name := claims.DisplayName
	if name == "" {
		name = claims.Username
	}
	return connector.Identity{
		UserID:            claims.UserID,
		Username:          name,
		PreferredUsername: claims.Username,
		Email:             standIn(claims.UserID),
		Groups:            claims.Groups,
	}, nil
}

func standIn(userID string) string {
	return userID + "@" + standInDomain
}

func (u *unified) Refresh(ctx context.Context, s connector.Scopes, ident connector.Identity) (connector.Identity, error) {
	userID, epoch, _, err := parseGrant(ident.ConnectorData)
	if err != nil || userID != ident.UserID {
		return connector.Identity{}, ErrGrantMalformed
	}
	next, err := u.identity(ctx, s, userID)
	if err != nil {
		return connector.Identity{}, err
	}
	user, err := u.deps.store.UserByID(ctx, userID)
	if err != nil {
		return connector.Identity{}, fmt.Errorf("idp: read credentials epoch: %w", err)
	}
	if user.CredentialsEpoch != epoch {
		return connector.Identity{}, ErrGrantRevoked
	}
	next.ConnectorData = ident.ConnectorData
	return next, nil
}

func encodeGrant(userID string, epoch int64, browser string) []byte {
	g := userID + "." + strconv.FormatInt(epoch, 10)
	if browser != "" {
		g += "." + browser
	}
	return []byte(g)
}

func parseGrant(data []byte) (string, int64, string, error) {
	if len(data) <= userIDLength+1 || data[userIDLength] != '.' {
		return "", 0, "", ErrGrantMalformed
	}
	rest, browser, found := strings.Cut(string(data[userIDLength+1:]), ".")
	if found && !storage.ValidBrowser(browser) {
		return "", 0, "", ErrGrantMalformed
	}
	epoch, err := strconv.ParseInt(rest, 10, 64)
	if err != nil || epoch < 0 {
		return "", 0, "", ErrGrantMalformed
	}
	return string(data[:userIDLength]), epoch, browser, nil
}