package store

import (
	"context"
	"strings"
)

type CacheAffinityObservationInput struct {
	CacheAffinityKey       string
	CacheAffinityPolicy    map[string]any
	PlatformID             string
	PlatformModelID        string
	ClientID               string
	ModelType              string
	InputTokens            int
	CachedInputTokens      int
	CachedInputTokensKnown bool
}

func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheAffinityObservationInput) error {
	input.CacheAffinityKey = strings.TrimSpace(input.CacheAffinityKey)
	input.ClientID = strings.TrimSpace(input.ClientID)
	input.ModelType = strings.TrimSpace(input.ModelType)
	if input.CacheAffinityKey == "" || input.ClientID == "" || input.ModelType == "" {
		return nil
	}
	if !input.CachedInputTokensKnown || !cacheAffinityPolicyEnabled(input.CacheAffinityPolicy, input.ModelType) {
		return nil
	}
	minInputTokens := cacheAffinityMinInputTokens(input.CacheAffinityPolicy)
	if input.InputTokens < minInputTokens || input.InputTokens <= 0 {
		return nil
	}
	if input.CachedInputTokens < 0 {
		input.CachedInputTokens = 0
	}
	if input.CachedInputTokens > input.InputTokens {
		input.CachedInputTokens = input.InputTokens
	}
	alpha := cacheAffinityEMAAlpha(input.CacheAffinityPolicy)
	hitRatio := float64(input.CachedInputTokens) / float64(input.InputTokens)
	_, err := s.pool.Exec(ctx, `
	INSERT INTO gateway_cache_affinity_stats (
	  client_id, cache_affinity_key, platform_id, platform_model_id, model_type,
	  request_count, input_tokens, cached_input_tokens, ema_hit_ratio, last_hit_ratio, last_observed_at
	)
	VALUES (
	  $1, $2, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, $5,
	  1, $6, $7, $8, $8, now()
	)
	ON CONFLICT (client_id, cache_affinity_key) DO UPDATE
	SET platform_id = COALESCE(EXCLUDED.platform_id, gateway_cache_affinity_stats.platform_id),
	    platform_model_id = COALESCE(EXCLUDED.platform_model_id, gateway_cache_affinity_stats.platform_model_id),
	    model_type = EXCLUDED.model_type,
	    request_count = gateway_cache_affinity_stats.request_count + 1,
	    input_tokens = gateway_cache_affinity_stats.input_tokens + EXCLUDED.input_tokens,
	    cached_input_tokens = gateway_cache_affinity_stats.cached_input_tokens + EXCLUDED.cached_input_tokens,
	    ema_hit_ratio = ($9::float8 * EXCLUDED.last_hit_ratio) + ((1 - $9::float8) * gateway_cache_affinity_stats.ema_hit_ratio),
	    last_hit_ratio = EXCLUDED.last_hit_ratio,
	    last_observed_at = now(),
	    updated_at = now()`,
		input.ClientID,
		input.CacheAffinityKey,
		input.PlatformID,
		input.PlatformModelID,
		input.ModelType,
		input.InputTokens,
		input.CachedInputTokens,
		hitRatio,
		alpha,
	)
	return err
}