package oauth

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"time"

	redispkg "github.com/ClareAI/astra-insight-mcp/internal/pkg/redis"
)

const tokenKeyPrefix = "astra:insight_mcp:wa_tokens:"

// TokenStore persists per-user (wa_id) OAuth tokens in Redis.
// Redis is the single source of truth — no local cache — so logout
// and token invalidation take effect immediately across all pods.
type TokenStore struct {
	redis *redispkg.Client
}

func NewTokenStore(rdb *redispkg.Client) *TokenStore {
	return &TokenStore{redis: rdb}
}

func (s *TokenStore) Get(waID string) *TokenData {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	val, err := s.redis.Get(ctx, tokenKeyPrefix+waID)
	if err != nil {
		slog.Warn("token_store: redis get failed", "wa_id", waID, "error", err)
		return nil
	}
	if val == "" {
		return nil
	}

	var tok TokenData
	if err := json.Unmarshal([]byte(val), &tok); err != nil {
		slog.Warn("token_store: unmarshal failed", "wa_id", waID, "error", err)
		return nil
	}
	return &tok
}

func (s *TokenStore) Set(waID string, tok *TokenData) error {
	data, err := json.Marshal(tok)
	if err != nil {
		return fmt.Errorf("marshal token: %w", err)
	}

	ttl := 24 * time.Hour
	if tok.ExpiresIn > 0 {
		ttl = time.Duration(tok.ExpiresIn) * time.Second
	}

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	if err := s.redis.Set(ctx, tokenKeyPrefix+waID, string(data), ttl); err != nil {
		slog.Error("token_store: redis set failed", "wa_id", waID, "error", err)
		return fmt.Errorf("redis set: %w", err)
	}
	return nil
}

func (s *TokenStore) Delete(waID string) {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	if err := s.redis.Del(ctx, tokenKeyPrefix+waID); err != nil {
		slog.Error("token_store: redis del failed", "wa_id", waID, "error", err)
	}
}

// RefreshIfNeeded checks whether the token is expired and refreshes it.
func (s *TokenStore) RefreshIfNeeded(waID string) *TokenData {
	tok := s.Get(waID)
	if tok == nil {
		return nil
	}
	if !tok.IsExpired() {
		return tok
	}
	if tok.RefreshToken == "" {
		return tok
	}
	newTok, err := RefreshAccessToken(tok)
	if err != nil {
		return tok
	}
	_ = s.Set(waID, newTok)
	return newTok
}
