package cache

import (
	"context"
	"crypto/sha256"
	"fmt"
	"time"

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

type Cache struct {
	client *redis.Client
	ttl    time.Duration
}

func New(addr string, ttlSeconds int) (*Cache, error) {
	client := redis.NewClient(&redis.Options{
		Addr: addr,
	})

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

	if err := client.Ping(ctx).Err(); err != nil {
		return nil, fmt.Errorf("connecting to Valkey: %w", err)
	}

	ttl := time.Duration(ttlSeconds) * time.Second
	if ttl == 0 {
		ttl = 1 * time.Hour
	}

	return &Cache{client: client, ttl: ttl}, nil
}

func (c *Cache) Get(ctx context.Context, model string, requestBody []byte) ([]byte, error) {
	key := c.cacheKey(model, requestBody)
	data, err := c.client.Get(ctx, key).Bytes()
	if err == redis.Nil {
		return nil, nil
	}
	return data, err
}

func (c *Cache) Set(ctx context.Context, model string, requestBody, responseBody []byte) error {
	key := c.cacheKey(model, requestBody)
	return c.client.Set(ctx, key, responseBody, c.ttl).Err()
}

func (c *Cache) Ping(ctx context.Context) error {
	return c.client.Ping(ctx).Err()
}

func (c *Cache) Close() error {
	return c.client.Close()
}

func (c *Cache) cacheKey(model string, requestBody []byte) string {
	h := sha256.New()
	h.Write([]byte(model))
	h.Write(requestBody)
	return fmt.Sprintf("llm-gw:%x", h.Sum(nil))
}