package auth

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"math/big"
	"slices"
	"strings"
	"sync"
	"time"
)

// DeviceFlow implements the OAuth 2.0 Device Authorization Grant
// (RFC 8628) in-process: a companion app requests a device + user code, the
// user approves it in a browser (after logging in), and the app polls until
// the flow mints an access token (a LocalJWT).
//
// State is held in memory: pending grants are short-lived (default 10 min)
// and minted tokens are stateless JWTs, so a backoffice restart only drops
// in-flight approvals — never issued tokens.
type DeviceFlow struct {
	mu       sync.Mutex
	pending  map[string]*deviceSession // deviceCode → session
	byCode   map[string]string         // userCode → deviceCode
	jwt      *LocalJWT
	baseURL  string        // issuer base for verification_uri
	ttl      time.Duration // device_code / user_code lifetime
	tokenTTL time.Duration // minted access-token lifetime
	interval time.Duration // client polling interval
	now      func() time.Time
}

type deviceSession struct {
	deviceCode string
	userCode   string
	expiresAt  time.Time
	requested  []string   // scopes the client asked for (RFC 8628 scope param)
	approved   *Principal // set when the user approves
}

// DeviceFlowOptions tune the flow. Zero values pick RFC-sensible defaults.
type DeviceFlowOptions struct {
	BaseURL  string        // e.g. "https://admin.example.com" for verification_uri
	TTL      time.Duration // device/user code lifetime (default 10m)
	TokenTTL time.Duration // access token lifetime (default 12h)
	Interval time.Duration // polling interval hint (default 5s)
}

// NewDeviceFlow constructs the flow. jwt must be enabled (non-empty secret)
// for Token() to work; Start/Approve still function.
func NewDeviceFlow(jwt *LocalJWT, opts DeviceFlowOptions) *DeviceFlow {
	if opts.TTL <= 0 {
		opts.TTL = 10 * time.Minute
	}
	if opts.TokenTTL <= 0 {
		opts.TokenTTL = 12 * time.Hour
	}
	if opts.Interval <= 0 {
		opts.Interval = 5 * time.Second
	}
	return &DeviceFlow{
		pending:  map[string]*deviceSession{},
		byCode:   map[string]string{},
		jwt:      jwt,
		baseURL:  strings.TrimSuffix(opts.BaseURL, "/"),
		ttl:      opts.TTL,
		tokenTTL: opts.TokenTTL,
		interval: opts.Interval,
		now:      time.Now,
	}
}

// Start begins a device authorization and returns the grant response the
// client sends to the user: device_code (for polling), user_code (typed by
// the user), verification_uri, verification_uri_complete (a direct link with
// the code pre-filled), expires_in (seconds) and the polling interval.
//
// requested lists the scopes the client wants on the minted token (e.g.
// ["mcp","read"] for a companion app). The approver's consent — plus the
// client having asked for them — is what grants them; Approve records the
// request on the session and Token mints with the intersection of the
// requested scopes and the approver's authorization.
func (f *DeviceFlow) Start(requested ...string) (deviceCode, userCode, verifyURI, verifyComplete string, expiresIn, interval int, err error) {
	if f == nil {
		return "", "", "", "", 0, 0, errors.New("auth: nil device flow")
	}
	dc, err := randomToken(24)
	if err != nil {
		return "", "", "", "", 0, 0, err
	}
	uc, err := genUserCode()
	if err != nil {
		return "", "", "", "", 0, 0, err
	}
	f.mu.Lock()
	f.pending[dc] = &deviceSession{deviceCode: dc, userCode: uc, expiresAt: f.now().Add(f.ttl), requested: append([]string(nil), requested...)}
	f.byCode[uc] = dc
	f.mu.Unlock()

	verify := f.baseURL
	if verify == "" {
		verify = "/device"
	} else {
		verify += "/device"
	}
	return dc, uc, verify, verify + "?code=" + uc, int(f.ttl / time.Second), int(f.interval / time.Second), nil
}

var (
	// ErrDeviceUnknown is returned by Approve for an unknown/expired user code.
	ErrDeviceUnknown = errors.New("auth: unknown or expired device code")
	// ErrAuthorizationPending is returned by Token while the user has not
	// approved yet (the client should poll again after the interval).
	ErrAuthorizationPending = errors.New("auth: authorization pending")
	// ErrDeviceExpired is returned by Token once the grant has expired.
	ErrDeviceExpired = errors.New("auth: device code expired")
)

// Approve binds the approving principal to the pending grant. The token's
// scopes are the client's REQUESTED scopes intersected with the approver's
// scopes; when the approver carries none (e.g. a browser admin session whose
// principal has no explicit scope list), the requested scopes pass through
// — consent from a logged-in administrator is the authorization.
//
// The approved principal's Subject still rides on the minted token (the
// audit trail says who approved), but it is not what grants scopes.
func (f *DeviceFlow) Approve(userCode string, approver *Principal) error {
	if f == nil {
		return ErrDeviceUnknown
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	dc, ok := f.byCode[userCode]
	if !ok {
		return ErrDeviceUnknown
	}
	s, ok := f.pending[dc]
	if !ok || f.now().After(s.expiresAt) {
		delete(f.byCode, userCode)
		return ErrDeviceUnknown
	}
	s.approved = approver
	return nil
}

// RequestedScopes returns the scopes a pending grant's client asked for.
// The approve UI shows them so the admin consents to exactly what the
// companion app will get. Returns (nil, false) for unknown/expired codes.
func (f *DeviceFlow) RequestedScopes(userCode string) ([]string, bool) {
	if f == nil {
		return nil, false
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	dc, ok := f.byCode[userCode]
	if !ok {
		return nil, false
	}
	s, ok := f.pending[dc]
	if !ok || f.now().After(s.expiresAt) {
		return nil, false
	}
	return append([]string(nil), s.requested...), true
}

// Token polls the grant: pending → ErrAuthorizationPending, expired →
// ErrDeviceExpired, approved → a minted access token (LocalJWT) for the
// approver. One-shot: a consumed grant is removed (tokens are stateless, so
// repeat polling would otherwise mint unlimited tokens from one approval).
func (f *DeviceFlow) Token(deviceCode string) (string, error) {
	if f == nil {
		return "", ErrDeviceUnknown
	}
	if f.jwt == nil || !f.jwt.Enabled() {
		return "", errors.New("auth: device tokens disabled (AUTH_TOKEN_SECRET not set)")
	}
	f.mu.Lock()
	s, ok := f.pending[deviceCode]
	if !ok {
		f.mu.Unlock()
		return "", ErrDeviceUnknown
	}
	if f.now().After(s.expiresAt) {
		delete(f.pending, deviceCode)
		delete(f.byCode, s.userCode)
		f.mu.Unlock()
		return "", ErrDeviceExpired
	}
	if s.approved == nil {
		f.mu.Unlock()
		return "", ErrAuthorizationPending
	}
	approver := s.approved
	delete(f.pending, deviceCode)
	delete(f.byCode, s.userCode)
	f.mu.Unlock()

	token, err := f.jwt.Mint(approver.Subject, grantScopes(s.requested, approver.Scopes), f.tokenTTL)
	if err != nil {
		return "", err
	}
	return token, nil
}

// grantScopes returns the intersection of the client's requested scopes and
// the approver's scopes. An approver with no explicit scopes (browser admin
// session) passes the requested scopes through — see Approve. An approver
// holding the "admin" scope may grant any requested scope (admin is the
// super-scope that implies the rest).
func grantScopes(requested, approver []string) []string {
	if len(requested) == 0 {
		return append([]string(nil), approver...)
	}
	if len(approver) == 0 {
		return append([]string(nil), requested...)
	}
	var out []string
	for _, r := range requested {
		if slices.Contains(approver, "admin") {
			out = append(out, r)
			continue
		}
		for _, a := range approver {
			if r == a {
				out = append(out, r)
				break
			}
		}
	}
	return out
}

// randomToken returns n random bytes as a URL-safe string.
func randomToken(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

// genUserCode generates a human-typable code in the GitHub style: two groups
// of four characters from an unambiguous charset (no 0/O, 1/I/L).
const userCodeAlphabet = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"

func genUserCode() (string, error) {
	b := make([]byte, 8)
	max := big.NewInt(int64(len(userCodeAlphabet)))
	for i := range b {
		n, err := rand.Int(rand.Reader, max)
		if err != nil {
			return "", err
		}
		b[i] = userCodeAlphabet[n.Int64()]
	}
	return string(b[:4]) + "-" + string(b[4:]), nil
}