package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

var ErrInvalidWSTicket = errors.New("invalid or expired websocket ticket")

type WSTicketClaims struct {
	UserID    uint   `json:"user_id"`
	AccountID uint   `json:"account_id"`
	Role      string `json:"role"`
	Provider  string `json:"provider"`
	ClientID  string `json:"client_id,omitempty"`
}

type wsTicketEntry struct {
	claims    WSTicketClaims
	expiresAt time.Time
}

// WSTicketStore exchanges a short-lived opaque ticket for user claims exactly once.
type WSTicketStore struct {
	rdb *redis.Client
	ttl time.Duration
	mu  sync.Mutex
	mem map[string]wsTicketEntry
}

func NewWSTicketStore(rdb *redis.Client, ttl time.Duration) *WSTicketStore {
	if ttl <= 0 {
		ttl = 30 * time.Second
	}
	return &WSTicketStore{rdb: rdb, ttl: ttl, mem: make(map[string]wsTicketEntry)}
}

func (s *WSTicketStore) Issue(ctx context.Context, claims WSTicketClaims) (string, error) {
	random := make([]byte, 32)
	if _, err := rand.Read(random); err != nil {
		return "", fmt.Errorf("generate websocket ticket: %w", err)
	}
	ticket := base64.RawURLEncoding.EncodeToString(random)
	key := wsTicketKey(ticket)
	if s.rdb == nil {
		s.mu.Lock()
		s.mem[key] = wsTicketEntry{claims: claims, expiresAt: time.Now().Add(s.ttl)}
		s.mu.Unlock()
		return ticket, nil
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		return "", err
	}
	if err := s.rdb.Set(ctx, key, payload, s.ttl).Err(); err != nil {
		return "", fmt.Errorf("store websocket ticket: %w", err)
	}
	return ticket, nil
}

func (s *WSTicketStore) Consume(ctx context.Context, ticket string) (*WSTicketClaims, error) {
	if ticket == "" {
		return nil, ErrInvalidWSTicket
	}
	key := wsTicketKey(ticket)
	if s.rdb == nil {
		s.mu.Lock()
		entry, ok := s.mem[key]
		delete(s.mem, key)
		s.mu.Unlock()
		if !ok || time.Now().After(entry.expiresAt) {
			return nil, ErrInvalidWSTicket
		}
		return &entry.claims, nil
	}
	payload, err := s.rdb.GetDel(ctx, key).Bytes()
	if errors.Is(err, redis.Nil) {
		return nil, ErrInvalidWSTicket
	}
	if err != nil {
		return nil, fmt.Errorf("consume websocket ticket: %w", err)
	}
	var claims WSTicketClaims
	if err := json.Unmarshal(payload, &claims); err != nil {
		return nil, ErrInvalidWSTicket
	}
	return &claims, nil
}

func wsTicketKey(ticket string) string {
	digest := sha256.Sum256([]byte(ticket))
	return "gochat:ws_ticket:" + hex.EncodeToString(digest[:])
}