package ai

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"
	"time"
)

const (
	ProviderXAI        = "xai"
	ProviderOpenCodeGo = "opencode-go"

	// ModelsCacheTTL is the fresh model-list window.
	ModelsCacheTTL = 5 * time.Minute
	// ModelsCacheStaleTTL is the maximum stale fallback window.
	ModelsCacheStaleTTL = 15 * time.Minute

	modelsRefreshLockTTL = 30 * time.Second
	modelsRefreshWait    = 25 * time.Millisecond
	modelsRefreshChecks  = 20
)

// ModelsCacheKey identifies one provider credential without containing the credential.
type ModelsCacheKey struct {
	Provider              string
	CredentialFingerprint string
}

// ModelsCacheValue is the JSON payload persisted by a ModelsCache.
type ModelsCacheValue struct {
	Models    []string `json:"models"`
	FetchedAt int64    `json:"fetchedAt"`
}

// ModelsCache stores model lists and coordinates refreshes across instances.
type ModelsCache interface {
	Get(ctx context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error)
	Set(ctx context.Context, key ModelsCacheKey, value ModelsCacheValue, ttl time.Duration) error
	AcquireRefresh(ctx context.Context, key ModelsCacheKey, token string, ttl time.Duration) (bool, error)
	ReleaseRefresh(ctx context.Context, key ModelsCacheKey, token string) error
}

// Registry holds the immutable provider client registry and a shared models cache.
type Registry struct {
	clients           map[string]Client
	cache             ModelsCache
	fingerprintSecret []byte
}

func NewRegistry(cache ModelsCache, fingerprintSecret string) *Registry {
	return newRegistry(cache, fingerprintSecret, map[string]Client{
		ProviderXAI:        NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
		ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
	})
}

// NewRegistryWithClients injects provider clients for focused tests.
func NewRegistryWithClients(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
	return newRegistry(cache, fingerprintSecret, clients)
}

func newRegistry(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
	immutableClients := make(map[string]Client, len(clients))
	for provider, client := range clients {
		immutableClients[NormalizeProvider(provider)] = client
	}
	return &Registry{clients: immutableClients, cache: cache, fingerprintSecret: []byte(fingerprintSecret)}
}

func NormalizeProvider(p string) string {
	switch strings.ToLower(strings.TrimSpace(p)) {
	case "opencode-go", "opencode", "opencode_go":
		return ProviderOpenCodeGo
	case "xai", "grok", "":
		return ProviderXAI
	default:
		return strings.ToLower(strings.TrimSpace(p))
	}
}

func (r *Registry) Client(provider string) (Client, error) {
	id := NormalizeProvider(provider)
	c, ok := r.clients[id]
	if !ok {
		return nil, fmt.Errorf("unsupported provider: %s (use xai or opencode-go)", provider)
	}
	return c, nil
}

// ListModelsCached fetches models via provider API; caches per provider+credential.
// fromCache indicates whether result was served from cache.
func (r *Registry) ListModelsCached(ctx context.Context, provider, apiKey string) (models []string, fromCache bool, err error) {
	id := NormalizeProvider(provider)
	if _, err := r.Client(id); err != nil {
		return nil, false, err
	}
	key := r.cacheKey(id, apiKey)
	now := time.Now().UTC()
	stale, hasStale := r.cached(ctx, key, now)
	if hasStale && cacheAge(now, stale.FetchedAt) <= ModelsCacheTTL {
		return cloneModels(stale.Models), true, nil
	}

	token, tokenErr := refreshToken()
	locked := false
	if tokenErr != nil {
		err = tokenErr
	} else if r.cache != nil {
		locked, err = r.cache.AcquireRefresh(ctx, key, token, modelsRefreshLockTTL)
		if err != nil {
			locked = false
		}
	}
	if locked {
		defer func() {
			releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second)
			defer cancel()
			_ = r.cache.ReleaseRefresh(releaseCtx, key, token)
		}()
	} else if r.cache != nil && err == nil {
		for range modelsRefreshChecks {
			select {
			case <-ctx.Done():
				return nil, false, ctx.Err()
			case <-time.After(modelsRefreshWait):
			}
			if current, ok := r.cached(ctx, key, time.Now().UTC()); ok && cacheAge(time.Now().UTC(), current.FetchedAt) <= ModelsCacheTTL {
				return cloneModels(current.Models), true, nil
			}
		}
	}

	c, err := r.Client(id)
	list, err := c.ListModels(ctx, apiKey)
	if err != nil {
		if hasStale {
			return cloneModels(stale.Models), true, nil
		}
		return nil, false, err
	}

	if r.cache != nil {
		_ = r.cache.Set(ctx, key, ModelsCacheValue{
			Models:    cloneModels(list),
			FetchedAt: time.Now().UTC().UnixNano(),
		}, ModelsCacheStaleTTL)
	}
	return list, false, nil
}

func (r *Registry) cached(ctx context.Context, key ModelsCacheKey, now time.Time) (ModelsCacheValue, bool) {
	if r.cache == nil {
		return ModelsCacheValue{}, false
	}
	value, ok, err := r.cache.Get(ctx, key)
	if err != nil || !ok || value.FetchedAt <= 0 {
		return ModelsCacheValue{}, false
	}
	if cacheAge(now, value.FetchedAt) > ModelsCacheStaleTTL {
		return ModelsCacheValue{}, false
	}
	return value, true
}

func (r *Registry) cacheKey(provider, apiKey string) ModelsCacheKey {
	mac := hmac.New(sha256.New, r.fingerprintSecret)
	_, _ = mac.Write([]byte(apiKey))
	return ModelsCacheKey{Provider: provider, CredentialFingerprint: hex.EncodeToString(mac.Sum(nil))}
}

func refreshToken() (string, error) {
	var token [16]byte
	if _, err := rand.Read(token[:]); err != nil {
		return "", err
	}
	return hex.EncodeToString(token[:]), nil
}

func cacheAge(now time.Time, fetchedAt int64) time.Duration {
	age := now.Sub(time.Unix(0, fetchedAt).UTC())
	if age < 0 {
		return 0
	}
	return age
}

func cloneModels(models []string) []string {
	return append([]string(nil), models...)
}

// FallbackModels when no key / provider error (never empty for UI).
// OpenCode Go 模型 id 見 https://opencode.ai/docs/go/（不是 OpenAI gpt-*）
func FallbackModels(provider string) []string {
	switch NormalizeProvider(provider) {
	case ProviderOpenCodeGo:
		return []string{
			"kimi-k2.6",
			"glm-5.1",
			"glm-5.2",
			"deepseek-v4-flash",
			"deepseek-v4-pro",
			"qwen3.7-plus",
			"minimax-m2.7",
			"mimo-v2.5",
		}
	default:
		return []string{"grok-3", "grok-3-mini", "grok-2"}
	}
}

// DefaultModel for provider when member settings model empty / invalid.
func DefaultModel(provider string) string {
	list := FallbackModels(provider)
	if len(list) == 0 {
		return "grok-3"
	}
	return list[0]
}

// SanitizeModel remaps invalid model ids for a provider (e.g. gpt-4o-mini on opencode-go).
func SanitizeModel(provider, model string) string {
	provider = NormalizeProvider(provider)
	model = strings.TrimSpace(model)
	if model == "" {
		return DefaultModel(provider)
	}
	if provider != ProviderOpenCodeGo {
		// xAI：允許 grok-*；若誤存 opencode 模型名，仍交由 API 驗證
		return model
	}
	// OpenCode Go：常見錯誤 — 從 xAI / OpenAI 帶過來的模型名
	lower := strings.ToLower(model)
	if strings.HasPrefix(lower, "gpt-") ||
		strings.HasPrefix(lower, "o1") ||
		strings.HasPrefix(lower, "o3") ||
		strings.HasPrefix(lower, "grok") ||
		strings.HasPrefix(lower, "claude") ||
		strings.HasPrefix(lower, "gemini") {
		return DefaultModel(provider)
	}
	return model
}