// Package redis supplies atomic shared middleware state in one Redis hash slot.
// Use a dedicated namespace and a noeviction Redis deployment for correctness.
package redis

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"git.yuntuo.hk/k2safe/go-omni/middleware"
	goredis "github.com/redis/go-redis/v9"
	"strings"
	"time"
)

type IdempotencyConfig struct {
	MaxEntries     int `json:"maxEntries"`
	MaxRecordBytes int `json:"maxRecordBytes"`
}
type IdempotencyStore struct {
	client    goredis.Scripter
	keys      []string
	cfg       IdempotencyConfig
	signature string
}

func namespace(prefix, kind string) (string, error) {
	if prefix == "" || len(prefix) > 256 || strings.ContainsAny(prefix, "{}\r\n") {
		return "", errors.New("middleware/redis: nonempty bounded namespace without braces required")
	}
	sum := sha256.Sum256([]byte(prefix))
	return "goOmni:{" + hex.EncodeToString(sum[:]) + "}:" + kind, nil
}
func NewIdempotencyStore(client goredis.Scripter, prefix string, cfg IdempotencyConfig) (*IdempotencyStore, error) {
	if client == nil || cfg.MaxEntries < 0 || cfg.MaxRecordBytes < 0 {
		return nil, errors.New("middleware/redis: invalid configuration")
	}
	if cfg.MaxEntries == 0 {
		cfg.MaxEntries = 128
	}
	if cfg.MaxRecordBytes == 0 {
		cfg.MaxRecordBytes = 128 << 10
	}
	if cfg.MaxEntries > 100000 || cfg.MaxRecordBytes < 512 || cfg.MaxRecordBytes > 1<<20 {
		return nil, errors.New("middleware/redis: capacity outside supported bounds")
	}
	base, err := namespace(prefix, "idempotency")
	if err != nil {
		return nil, err
	}
	return &IdempotencyStore{client: client, keys: []string{base + ":records", base + ":expiry", base + ":config"}, cfg: cfg, signature: fmt.Sprintf("%d:%d", cfg.MaxEntries, cfg.MaxRecordBytes)}, nil
}

var claim = goredis.NewScript(`
local config=redis.call('GET',KEYS[3])
if config and config~=ARGV[5] then return redis.error_reply('configuration mismatch') end
if not config then redis.call('SET',KEYS[3],ARGV[5]) end
local clock=redis.call('TIME');local now=clock[1]*1000+math.floor(clock[2]/1000)
local expired=redis.call('ZRANGEBYSCORE',KEYS[2],'-inf',now,'LIMIT',0,128)
for _,key in ipairs(expired) do redis.call('HDEL',KEYS[1],key);redis.call('ZREM',KEYS[2],key) end
local expiry=redis.call('ZSCORE',KEYS[2],ARGV[1]);if expiry and tonumber(expiry)<=now then redis.call('HDEL',KEYS[1],ARGV[1]);redis.call('ZREM',KEYS[2],ARGV[1]) end
local value=redis.call('HGET',KEYS[1],ARGV[1]);if value then return {0,value} end
if redis.call('HLEN',KEYS[1])>=tonumber(ARGV[4]) then return redis.error_reply('idempotency capacity reached') end
value=cjson.encode({fingerprint=ARGV[2],owner=ARGV[3],state='running'})
redis.call('HSET',KEYS[1],ARGV[1],value);return {1,value}
`)

func (s *IdempotencyStore) Claim(ctx context.Context, key, fp, owner string) (middleware.IdempotencyRecord, bool, error) {
	if len(key) != 64 || len(fp) != 64 || len(owner) != 32 {
		return middleware.IdempotencyRecord{}, false, errors.New("middleware/redis: invalid claim")
	}
	result, err := claim.Run(ctx, s.client, s.keys, key, fp, owner, s.cfg.MaxEntries, s.signature).Slice()
	if err != nil {
		return middleware.IdempotencyRecord{}, false, err
	}
	if len(result) != 2 {
		return middleware.IdempotencyRecord{}, false, errors.New("middleware/redis: invalid claim result")
	}
	value, ok := result[1].(string)
	if !ok || len(value) > s.cfg.MaxRecordBytes {
		return middleware.IdempotencyRecord{}, false, errors.New("middleware/redis: invalid record")
	}
	var record middleware.IdempotencyRecord
	err = json.Unmarshal([]byte(value), &record)
	return record, result[0] == int64(1), err
}

var complete = goredis.NewScript(`
local current=redis.call('HGET',KEYS[1],ARGV[1]);if not current then return redis.error_reply('claim missing') end
local value=cjson.decode(current)
if value.owner~=ARGV[2] or value.state~='running' then return redis.error_reply('claim owner/state changed') end
local response=ARGV[3]
if response~='' then value.state='completed';value.response=cjson.decode(response) else value.state='failed' end
local encoded=cjson.encode(value)
if string.len(encoded)>tonumber(ARGV[5]) then value.state='failed';value.response=nil;encoded=cjson.encode(value) end
local clock=redis.call('TIME');local now=clock[1]*1000+math.floor(clock[2]/1000)
redis.call('HSET',KEYS[1],ARGV[1],encoded);redis.call('ZADD',KEYS[2],now+tonumber(ARGV[4]),ARGV[1]);return 1
`)

func (s *IdempotencyStore) Complete(ctx context.Context, key, owner string, response *middleware.StoredResponse, ttl time.Duration) error {
	if ttl <= 0 || ttl > 30*24*time.Hour {
		return errors.New("middleware/redis: TTL must be positive and <=30d")
	}
	var encoded []byte
	var err error
	if response != nil {
		encoded, err = json.Marshal(response)
		if err != nil {
			return err
		}
		if len(encoded) > s.cfg.MaxRecordBytes {
			encoded = nil
		}
	}
	return complete.Run(ctx, s.client, s.keys, key, owner, string(encoded), max(1, ttl.Milliseconds()), s.cfg.MaxRecordBytes).Err()
}

// Resolve removes only a matching non-replayable claim after the operator has
// reconciled the business outcome. Never invoke automatically on a timeout.
func (s *IdempotencyStore) Resolve(ctx context.Context, key, owner string) error {
	return goredis.NewScript(`local v=redis.call('HGET',KEYS[1],ARGV[1]);if not v then return 0 end;local r=cjson.decode(v);if r.owner~=ARGV[2] or r.state=='completed' then return redis.error_reply('claim cannot be resolved') end;redis.call('HDEL',KEYS[1],ARGV[1]);redis.call('ZREM',KEYS[2],ARGV[1]);return 1`).Run(ctx, s.client, s.keys, key, owner).Err()
}

type limiter struct {
	client    goredis.Scripter
	keys      []string
	cfg       middleware.RateLimitConfig
	signature string
}

func NewRateLimiter(client goredis.Scripter, prefix string, cfg middleware.RateLimitConfig) (*middleware.RateLimiter, error) {
	if client == nil {
		return nil, errors.New("middleware/redis: client required")
	}
	if cfg.MaxKeys == 0 {
		cfg.MaxKeys = 10000
	}
	if cfg.IdleTTL == 0 {
		cfg.IdleTTL = 24 * time.Hour
	}
	if _, err := middleware.NewRateLimiter(cfg); err != nil {
		return nil, err
	}
	base, err := namespace(prefix, "rate")
	if err != nil {
		return nil, err
	}
	backend := &limiter{client: client, keys: []string{base + ":buckets", base + ":seen", base + ":config"}, cfg: cfg, signature: fmt.Sprintf("%g:%d:%d:%d", cfg.Rate, cfg.Burst, cfg.MaxKeys, cfg.IdleTTL.Milliseconds())}
	cfg.Backend = backend
	return middleware.NewRateLimiter(cfg)
}

var take = goredis.NewScript(`
local config=redis.call('GET',KEYS[3]);if config and config~=ARGV[6] then return redis.error_reply('configuration mismatch') end
if not config then redis.call('SET',KEYS[3],ARGV[6]) end
local clock=redis.call('TIME');local now=clock[1]*1000+math.floor(clock[2]/1000)
local expired=redis.call('ZRANGEBYSCORE',KEYS[2],'-inf',now-tonumber(ARGV[5]),'LIMIT',0,128)
for _,key in ipairs(expired) do redis.call('HDEL',KEYS[1],key);redis.call('ZREM',KEYS[2],key) end
local value=redis.call('HGET',KEYS[1],ARGV[1]);local tokens=tonumber(ARGV[3])
if value then local b=cjson.decode(value);tokens=math.min(tokens,b.tokens+math.max(0,now-b.time)*tonumber(ARGV[2])/1000)
elseif redis.call('HLEN',KEYS[1])>=tonumber(ARGV[4]) then return -1 end
local delay=0;if tokens>=1 then tokens=tokens-1 else delay=math.ceil((1-tokens)*1000/tonumber(ARGV[2])) end
redis.call('HSET',KEYS[1],ARGV[1],cjson.encode({tokens=tokens,time=now}));redis.call('ZADD',KEYS[2],now,ARGV[1]);return delay
`)

func (l *limiter) Take(ctx context.Context, key string) (time.Duration, error) {
	sum := sha256.Sum256([]byte(key))
	delay, err := take.Run(ctx, l.client, l.keys, hex.EncodeToString(sum[:]), l.cfg.Rate, l.cfg.Burst, l.cfg.MaxKeys, l.cfg.IdleTTL.Milliseconds(), l.signature).Int64()
	if err != nil {
		return 0, err
	}
	if delay < 0 {
		return 0, middleware.ErrLimitCapacity
	}
	return time.Duration(delay) * time.Millisecond, nil
}