package limits

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"math"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"aigw/internal/domain"

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

var (
	ErrRequestsExceeded = errors.New("requests per minute limit exceeded")
	ErrTokensExceeded   = errors.New("tokens per minute limit exceeded")
	ErrConcurrencyLimit = errors.New("concurrent request limit exceeded")
)

type Lease interface{ Release() }

type Limiter struct {
	redis            *redis.Client
	prefix           string
	defaultMaxOutput int64
	logger           *slog.Logger
	policies         atomic.Pointer[policySnapshot]
	localMu          sync.Mutex
	local            map[string]*localWindow
	redisUp          atomic.Bool
	redisFailureSeen atomic.Bool
	redisRetryAt     atomic.Int64
}

type policySnapshot struct{ policies map[string]domain.LimitPolicy }

type localWindow struct {
	minute     int64
	requests   int64
	tokens     int64
	concurrent int64
}

type localLease struct {
	limiter  *Limiter
	project  string
	released atomic.Bool
}
type redisLease struct {
	limiter  *Limiter
	key      string
	released atomic.Bool
}

const (
	redisCommandTimeout = 250 * time.Millisecond
	redisRetryCooldown  = 5 * time.Second
)

const acquireScript = `
local project_req = tonumber(ARGV[1])
local project_tok = tonumber(ARGV[2])
local project_conc = tonumber(ARGV[3])
local key_req = tonumber(ARGV[4])
local key_tok = tonumber(ARGV[5])
local estimate = tonumber(ARGV[6])
local ttl = tonumber(ARGV[7])
local pr = 0
local pt = 0
local c = 0
local kr = 0
local kt = 0
if project_req > 0 then pr = redis.call('INCR', KEYS[1]); if pr == 1 then redis.call('PEXPIRE', KEYS[1], ttl) end end
if project_tok > 0 then pt = redis.call('INCRBY', KEYS[2], estimate); if pt == estimate then redis.call('PEXPIRE', KEYS[2], ttl) end end
if project_conc > 0 then c = redis.call('INCR', KEYS[3]); redis.call('PEXPIRE', KEYS[3], 3600000) end
if key_req > 0 then kr = redis.call('INCR', KEYS[4]); if kr == 1 then redis.call('PEXPIRE', KEYS[4], ttl) end end
if key_tok > 0 then kt = redis.call('INCRBY', KEYS[5], estimate); if kt == estimate then redis.call('PEXPIRE', KEYS[5], ttl) end end
local reason = 0
if project_req > 0 and pr > project_req then reason = 1
elseif project_tok > 0 and pt > project_tok then reason = 2
elseif project_conc > 0 and c > project_conc then reason = 3
elseif key_req > 0 and kr > key_req then reason = 4
elseif key_tok > 0 and kt > key_tok then reason = 5 end
if reason > 0 then
  if project_req > 0 then redis.call('DECR', KEYS[1]) end
  if project_tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end
  if project_conc > 0 then redis.call('DECR', KEYS[3]) end
  if key_req > 0 then redis.call('DECR', KEYS[4]) end
  if key_tok > 0 then redis.call('DECRBY', KEYS[5], estimate) end
  return {0,reason}
end
return {1,0}
`

const releaseScript = `
local value = tonumber(redis.call('GET', KEYS[1]) or '0')
if value > 0 then redis.call('DECR', KEYS[1]) end
return value
`

func New(redisURL, prefix string, defaultMaxOutput int64, logger *slog.Logger) *Limiter {
	if logger == nil {
		logger = slog.Default()
	}
	if prefix == "" {
		prefix = "aigw:limits"
	}
	l := &Limiter{prefix: strings.TrimRight(prefix, ":"), defaultMaxOutput: defaultMaxOutput, logger: logger, local: make(map[string]*localWindow)}
	if strings.TrimSpace(redisURL) != "" {
		if options, err := redis.ParseURL(redisURL); err == nil {
			options.MaxRetries = -1
			options.DialerRetries = 1
			options.DialTimeout = redisCommandTimeout
			options.ReadTimeout = redisCommandTimeout
			options.WriteTimeout = redisCommandTimeout
			options.PoolTimeout = redisCommandTimeout
			l.redis = redis.NewClient(options)
		} else {
			logger.Warn("limits_redis_config_invalid", "error", err, "fallback", "local")
		}
	}
	l.ReplacePolicies(nil)
	return l
}

func (l *Limiter) Close() error {
	if l.redis == nil {
		return nil
	}
	return l.redis.Close()
}

func (l *Limiter) ReplacePolicies(policies []domain.LimitPolicy) {
	copyMap := make(map[string]domain.LimitPolicy, len(policies))
	for _, policy := range policies {
		if policy.ProjectID != "" {
			copyMap[policy.ProjectID] = policy
		}
	}
	l.policies.Store(&policySnapshot{policies: copyMap})
}

func (l *Limiter) Policy(projectID string) (domain.LimitPolicy, bool) {
	snapshot := l.policies.Load()
	if snapshot == nil {
		return domain.LimitPolicy{}, false
	}
	policy, ok := snapshot.policies[projectID]
	return policy, ok
}

func (l *Limiter) Acquire(ctx context.Context, principal domain.Principal, body []byte) (Lease, error) {
	policy, _ := l.Policy(principal.ProjectID)
	if policy.RequestsPerMinute == 0 && policy.TokensPerMinute == 0 && policy.Concurrent == 0 && principal.RequestsPerMinute == 0 && principal.TokensPerMinute == 0 {
		return noopLease{}, nil
	}
	estimate := EstimateTokens(body, l.defaultMaxOutput)
	minute := time.Now().Unix() / 60
	if l.redis != nil && time.Now().UnixNano() >= l.redisRetryAt.Load() {
		redisContext, cancel := context.WithTimeout(ctx, redisCommandTimeout)
		lease, err := l.acquireRedis(redisContext, principal, minute, policy, estimate)
		cancel()
		if err == nil {
			return lease, nil
		}
		if errors.Is(err, ErrRequestsExceeded) || errors.Is(err, ErrTokensExceeded) || errors.Is(err, ErrConcurrencyLimit) {
			l.markRedis(true, nil)
			return nil, err
		}
		l.markRedis(false, err)
	}
	return l.acquireLocal(principal, minute, policy, estimate)
}

func EstimateTokens(body []byte, defaultMaxOutput int64) int64 {
	input := int64((len(body) + 3) / 4)
	if input < 1 {
		input = 1
	}
	maxOutput := defaultMaxOutput
	var limits struct {
		MaxTokens           int64 `json:"max_tokens"`
		MaxCompletionTokens int64 `json:"max_completion_tokens"`
		MaxOutputTokens     int64 `json:"max_output_tokens"`
	}
	if json.Unmarshal(body, &limits) == nil {
		explicit := int64(0)
		for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} {
			if value > explicit {
				explicit = value
			}
		}
		if explicit > 0 {
			maxOutput = explicit
		}
	}
	if maxOutput < 0 {
		maxOutput = 0
	}
	if input > math.MaxInt64-maxOutput {
		return math.MaxInt64
	}
	return input + maxOutput
}

func (l *Limiter) acquireRedis(ctx context.Context, principal domain.Principal, minute int64, policy domain.LimitPolicy, estimate int64) (Lease, error) {
	projectBase := l.prefix + ":project:" + principal.ProjectID + ":" + strconv.FormatInt(minute, 10)
	keyBase := l.prefix + ":key:" + principal.KeyID + ":" + strconv.FormatInt(minute, 10)
	concurrencyKey := l.prefix + ":project:" + principal.ProjectID + ":concurrent"
	keys := []string{projectBase + ":requests", projectBase + ":tokens", concurrencyKey, keyBase + ":requests", keyBase + ":tokens"}
	values, err := l.redis.Eval(ctx, acquireScript, keys, policy.RequestsPerMinute, policy.TokensPerMinute, policy.Concurrent,
		principal.RequestsPerMinute, principal.TokensPerMinute, estimate, 125000).Result()
	if err != nil {
		return nil, err
	}
	items, ok := values.([]any)
	if !ok || len(items) < 2 {
		return nil, errors.New("invalid limits Redis response")
	}
	allowed, _ := toInt64(items[0])
	reason, _ := toInt64(items[1])
	if allowed == 0 {
		switch reason {
		case 1, 4:
			return nil, ErrRequestsExceeded
		case 2, 5:
			return nil, ErrTokensExceeded
		default:
			return nil, ErrConcurrencyLimit
		}
	}
	l.markRedis(true, nil)
	return &redisLease{limiter: l, key: concurrencyKey}, nil
}

func (l *Limiter) acquireLocal(principal domain.Principal, minute int64, policy domain.LimitPolicy, estimate int64) (Lease, error) {
	l.localMu.Lock()
	defer l.localMu.Unlock()
	for key, value := range l.local {
		if value.minute < minute-2 && value.concurrent == 0 {
			delete(l.local, key)
		}
	}
	projectKey := "project:" + principal.ProjectID
	keyKey := "key:" + principal.KeyID
	projectWindow := l.localWindow(projectKey, minute)
	keyWindow := l.localWindow(keyKey, minute)
	if policy.RequestsPerMinute > 0 && projectWindow.requests >= policy.RequestsPerMinute {
		return nil, ErrRequestsExceeded
	}
	if policy.TokensPerMinute > 0 && (estimate > policy.TokensPerMinute || projectWindow.tokens > policy.TokensPerMinute-estimate) {
		return nil, ErrTokensExceeded
	}
	if policy.Concurrent > 0 && projectWindow.concurrent >= policy.Concurrent {
		return nil, ErrConcurrencyLimit
	}
	if principal.RequestsPerMinute > 0 && keyWindow.requests >= principal.RequestsPerMinute {
		return nil, ErrRequestsExceeded
	}
	if principal.TokensPerMinute > 0 && (estimate > principal.TokensPerMinute || keyWindow.tokens > principal.TokensPerMinute-estimate) {
		return nil, ErrTokensExceeded
	}
	projectWindow.requests++
	projectWindow.tokens += estimate
	projectWindow.concurrent++
	keyWindow.requests++
	keyWindow.tokens += estimate
	return &localLease{limiter: l, project: projectKey}, nil
}

func (l *Limiter) localWindow(key string, minute int64) *localWindow {
	window := l.local[key]
	if window == nil {
		window = &localWindow{minute: minute}
		l.local[key] = window
	} else if window.minute != minute {
		window.minute = minute
		window.requests = 0
		window.tokens = 0
	}
	return window
}

func (l *Limiter) releaseLocal(project string) {
	l.localMu.Lock()
	defer l.localMu.Unlock()
	if value := l.local[project]; value != nil && value.concurrent > 0 {
		value.concurrent--
	}
}

func (l *Limiter) releaseRedis(ctx context.Context, key string) {
	if l.redis == nil {
		return
	}
	if _, err := l.redis.Eval(ctx, releaseScript, []string{key}).Result(); err != nil {
		l.markRedis(false, err)
	}
}

func (x *localLease) Release() {
	if x.released.CompareAndSwap(false, true) {
		x.limiter.releaseLocal(x.project)
	}
}
func (x *redisLease) Release() {
	if x.released.CompareAndSwap(false, true) {
		ctx, cancel := context.WithTimeout(context.Background(), redisCommandTimeout)
		defer cancel()
		x.limiter.releaseRedis(ctx, x.key)
	}
}

type noopLease struct{}

func (noopLease) Release() {}

func (l *Limiter) markRedis(healthy bool, err error) {
	l.redisUp.Store(healthy)
	if healthy {
		l.redisRetryAt.Store(0)
		if l.redisFailureSeen.Swap(false) {
			l.logger.Info("limits_redis_recovered")
		}
		return
	}
	l.redisRetryAt.Store(time.Now().Add(redisRetryCooldown).UnixNano())
	if l.redisFailureSeen.CompareAndSwap(false, true) {
		l.logger.Warn("limits_redis_unavailable", "error", err, "fallback", "local", "retry_after", redisRetryCooldown)
	}
}

func toInt64(value any) (int64, bool) {
	switch v := value.(type) {
	case int64:
		return v, true
	case string:
		n, e := strconv.ParseInt(v, 10, 64)
		return n, e == nil
	case []byte:
		n, e := strconv.ParseInt(string(v), 10, 64)
		return n, e == nil
	default:
		return 0, false
	}
}