package runtime

import (
	"crypto/sha256"
	"fmt"
	"math/rand"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/apr/types"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
)

const (
	DefaultProvider     = "copilot"
	DefaultSyncInterval = 5 * time.Minute
)

var providerRedirects = map[string]string{"glm": "zhipu"}

// Dispatcher is the APR-style model dispatcher without AgentTeam auth/config globals.
type Dispatcher struct {
	mutex          sync.RWMutex
	providers      map[string]types.LLMProvider
	modelProviders map[string][]string
	routes         []ModelRoute
	replacements   []ModelReplacement
	syncInterval   time.Duration
	stopCh         chan struct{}
	syncDone       chan struct{}
	stopOnce       sync.Once
	randomSource   *rand.Rand
}

// NewDispatcher builds one dispatcher instance.
func NewDispatcher() *Dispatcher {
	dispatcher := &Dispatcher{
		providers:      make(map[string]types.LLMProvider),
		modelProviders: make(map[string][]string),
		routes:         []ModelRoute{},
		syncInterval:   DefaultSyncInterval,
		stopCh:         make(chan struct{}),
		syncDone:       make(chan struct{}),
		randomSource:   rand.New(rand.NewSource(time.Now().UnixNano())),
	}
	logging.Go(nil, func() { dispatcher.syncLoop() })
	return dispatcher
}

// AddProvider registers one provider and syncs its models into the direct lookup map.
func (dispatcher *Dispatcher) AddProvider(provider types.LLMProvider) {
	if provider == nil {
		return
	}
	name := strings.ToLower(strings.TrimSpace(provider.Name()))
	if name == "" {
		return
	}

	dispatcher.mutex.Lock()
	defer dispatcher.mutex.Unlock()
	dispatcher.providers[name] = provider
	dispatcher.syncProviderModels(name, provider)
}

// Reset clears all registered providers, routes, replacements, and model indexes.
func (dispatcher *Dispatcher) Reset() {
	dispatcher.mutex.Lock()
	defer dispatcher.mutex.Unlock()

	dispatcher.providers = make(map[string]types.LLMProvider)
	dispatcher.modelProviders = make(map[string][]string)
	dispatcher.routes = nil
	dispatcher.replacements = nil
}

// SetRoutes replaces route/group configuration.
func (dispatcher *Dispatcher) SetRoutes(routes []ModelRoute) {
	dispatcher.mutex.Lock()
	defer dispatcher.mutex.Unlock()

	dispatcher.routes = make([]ModelRoute, len(routes))
	for index, route := range routes {
		dispatcher.routes[index] = ModelRoute{
			ID:        route.ID,
			Vendor:    strings.ToLower(strings.TrimSpace(route.Vendor)),
			Models:    normalizeStrings(route.Models),
			Providers: normalizeStrings(route.Providers),
			Groups:    normalizeGroups(route.Groups),
		}
	}
}

// SetModelReplacements replaces model replacement rules while preserving counters
// for rules with stable IDs.
func (dispatcher *Dispatcher) SetModelReplacements(replacements []ModelReplacement) {
	dispatcher.mutex.Lock()
	defer dispatcher.mutex.Unlock()

	oldMap := make(map[string]ModelReplacement, len(dispatcher.replacements))
	for _, replacement := range dispatcher.replacements {
		oldMap[replacement.ID] = replacement
	}

	dispatcher.replacements = make([]ModelReplacement, len(replacements))
	for index, replacement := range replacements {
		replacement.RequestCount = 0
		replacement.HitCount = 0
		if strings.TrimSpace(replacement.ID) == "" {
			replacement.ID = computeReplacementID(&replacement)
		}
		if old, ok := oldMap[replacement.ID]; ok {
			replacement.RequestCount = old.RequestCount
			replacement.HitCount = old.HitCount
		}
		dispatcher.replacements[index] = replacement
	}
}

// ReplacementStats returns a snapshot of the current replacement rules with their stats.
func (dispatcher *Dispatcher) ReplacementStats() []ModelReplacement {
	dispatcher.mutex.RLock()
	defer dispatcher.mutex.RUnlock()

	result := make([]ModelReplacement, len(dispatcher.replacements))
	for index, replacement := range dispatcher.replacements {
		result[index] = ModelReplacement{
			ID:           replacement.ID,
			APIKey:       replacement.APIKey,
			Model:        replacement.Model,
			TargetModel:  replacement.TargetModel,
			Probability:  replacement.Probability,
			RequestCount: replacement.RequestCount,
			HitCount:     replacement.HitCount,
		}
	}
	return result
}

// ProviderModels returns all registered providers and their models.
func (dispatcher *Dispatcher) ProviderModels() map[string][]types.Model {
	dispatcher.mutex.RLock()
	defer dispatcher.mutex.RUnlock()

	result := make(map[string][]types.Model, len(dispatcher.providers))
	for name, provider := range dispatcher.providers {
		result[name] = provider.Models()
	}
	return result
}

// Models returns visible models grouped by vendor, filtered by routes/groups.
func (dispatcher *Dispatcher) Models() map[string][]types.Model {
	dispatcher.mutex.RLock()
	defer dispatcher.mutex.RUnlock()

	result := make(map[string][]types.Model)
	seen := make(map[string]map[string]bool)

	for _, route := range dispatcher.routes {
		for _, group := range route.EffectiveGroups() {
			if len(group.Models) == 0 {
				continue
			}
			allModelsLower := make(map[string]types.Model)
			for _, providerName := range group.Providers {
				provider, ok := dispatcher.providers[strings.ToLower(strings.TrimSpace(providerName))]
				if !ok {
					continue
				}
				for _, model := range provider.Models() {
					allModelsLower[strings.ToLower(model.ID)] = model
				}
			}
			if seen[route.Vendor] == nil {
				seen[route.Vendor] = make(map[string]bool)
			}
			for _, name := range group.Models {
				if model, ok := allModelsLower[strings.ToLower(name)]; ok {
					if !seen[route.Vendor][model.ID] {
						seen[route.Vendor][model.ID] = true
						result[route.Vendor] = append(result[route.Vendor], model)
					}
				}
			}
		}
	}

	return result
}

// Lookup resolves the provider and normalized model name for one request model.
func (dispatcher *Dispatcher) Lookup(apiKey string, model string) (types.LLMProvider, string) {
	dispatcher.mutex.RLock()
	defer dispatcher.mutex.RUnlock()

	model = dispatcher.applyModelReplacement(apiKey, model)
	return dispatcher.lookup(model, nil)
}

// LookupAmongProviders resolves within one allowed provider subset.
func (dispatcher *Dispatcher) LookupAmongProviders(apiKey string, model string, allowedProviders []string) (types.LLMProvider, string) {
	dispatcher.mutex.RLock()
	defer dispatcher.mutex.RUnlock()

	model = dispatcher.applyModelReplacement(apiKey, model)
	allowed := make(map[string]struct{}, len(allowedProviders))
	for _, providerName := range allowedProviders {
		trimmed := strings.ToLower(strings.TrimSpace(providerName))
		if trimmed == "" {
			continue
		}
		allowed[trimmed] = struct{}{}
	}
	return dispatcher.lookup(model, allowed)
}

func (dispatcher *Dispatcher) lookup(model string, allowed map[string]struct{}) (types.LLMProvider, string) {
	parts := strings.SplitN(strings.TrimSpace(model), "/", 2)
	if len(parts) < 2 {
		return dispatcher.lookupDirect(model, allowed)
	}

	prefix := strings.ToLower(parts[0])
	modelName := strings.ToLower(parts[1])

	var providers []types.LLMProvider
	seen := make(map[string]bool)

	for _, route := range dispatcher.routes {
		if route.Vendor != prefix {
			continue
		}
		for _, group := range route.EffectiveGroups() {
			if len(group.Models) > 0 && !containsFold(group.Models, modelName) {
				continue
			}
			for _, providerName := range group.Providers {
				normalizedProviderName := strings.ToLower(strings.TrimSpace(providerName))
				if seen[normalizedProviderName] {
					continue
				}
				if !allowedContains(allowed, normalizedProviderName) {
					continue
				}
				provider, ok := dispatcher.providers[normalizedProviderName]
				if !ok {
					continue
				}
				seen[normalizedProviderName] = true
				providers = append(providers, provider)
			}
		}
	}

	if provider, ok := dispatcher.providers[prefix]; ok && !seen[prefix] && allowedContains(allowed, prefix) {
		seen[prefix] = true
		providers = append(providers, provider)
	}

	if redirect, ok := providerRedirects[prefix]; ok && !seen[redirect] && allowedContains(allowed, redirect) {
		if provider, ok := dispatcher.providers[redirect]; ok {
			seen[redirect] = true
			providers = append(providers, provider)
		}
	}

	if len(providers) > 0 {
		return providers[dispatcher.randomSource.Intn(len(providers))], modelName
	}

	return nil, ""
}

func (dispatcher *Dispatcher) lookupDirect(model string, allowed map[string]struct{}) (types.LLMProvider, string) {
	modelName := strings.ToLower(strings.TrimSpace(model))
	if modelName == "" {
		return nil, ""
	}

	if providers, ok := dispatcher.modelProviders[modelName]; ok && len(providers) > 0 {
		candidates := make([]string, 0, len(providers))
		for _, providerName := range providers {
			if allowedContains(allowed, providerName) {
				candidates = append(candidates, providerName)
			}
		}
		if len(candidates) > 0 {
			providerName := candidates[dispatcher.randomSource.Intn(len(candidates))]
			if provider, ok := dispatcher.providers[providerName]; ok {
				return provider, modelName
			}
		}
	}

	if allowedContains(allowed, DefaultProvider) {
		if provider, ok := dispatcher.providers[DefaultProvider]; ok {
			return provider, model
		}
	}

	return nil, model
}

func (dispatcher *Dispatcher) applyModelReplacement(apiKey string, model string) string {
	if apiKey == "" || len(dispatcher.replacements) == 0 {
		return model
	}
	for index := range dispatcher.replacements {
		replacement := &dispatcher.replacements[index]
		if replacement.APIKey != apiKey || replacement.Probability <= 0 {
			continue
		}
		matched := replacement.Model == "*"
		if !matched {
			matched, _ = filepath.Match(strings.ToLower(replacement.Model), strings.ToLower(model))
		}
		if !matched {
			continue
		}
		atomic.AddUint64(&replacement.RequestCount, 1)
		if replacement.Probability >= 1 || dispatcher.randomSource.Float64() < replacement.Probability {
			atomic.AddUint64(&replacement.HitCount, 1)
			return replacement.TargetModel
		}
		return model
	}
	return model
}

// Sync rebuilds the direct model lookup index from current providers.
func (dispatcher *Dispatcher) Sync() {
	dispatcher.mutex.Lock()
	defer dispatcher.mutex.Unlock()

	dispatcher.modelProviders = make(map[string][]string)
	for name, provider := range dispatcher.providers {
		dispatcher.syncProviderModels(name, provider)
	}
}

func (dispatcher *Dispatcher) syncLoop() {
	ticker := time.NewTicker(dispatcher.syncInterval)
	defer ticker.Stop()
	defer close(dispatcher.syncDone)

	dispatcher.Sync()

	for {
		select {
		case <-ticker.C:
			dispatcher.Sync()
		case <-dispatcher.stopCh:
			return
		}
	}
}

// Stop stops the background sync goroutine.
func (dispatcher *Dispatcher) Stop() {
	if dispatcher == nil {
		return
	}
	dispatcher.stopOnce.Do(func() {
		close(dispatcher.stopCh)
	})
	<-dispatcher.syncDone
}

func (dispatcher *Dispatcher) syncProviderModels(name string, provider types.LLMProvider) {
	for _, model := range provider.Models() {
		modelID := strings.ToLower(strings.TrimSpace(model.ID))
		if modelID == "" {
			continue
		}
		dispatcher.modelProviders[modelID] = appendIfMissing(dispatcher.modelProviders[modelID], name)
	}
}

func computeReplacementID(replacement *ModelReplacement) string {
	hash := sha256.Sum256([]byte(fmt.Sprintf("%s:%s:%s:%v", replacement.APIKey, replacement.Model, replacement.TargetModel, replacement.Probability)))
	return fmt.Sprintf("%x", hash[:8])
}

func normalizeGroups(groups []RouteGroup) []RouteGroup {
	result := make([]RouteGroup, 0, len(groups))
	for _, group := range groups {
		result = append(result, RouteGroup{
			ID:        group.ID,
			Models:    normalizeStrings(group.Models),
			Providers: normalizeStrings(group.Providers),
		})
	}
	return result
}

func normalizeStrings(items []string) []string {
	result := make([]string, 0, len(items))
	for _, item := range items {
		trimmed := strings.ToLower(strings.TrimSpace(item))
		if trimmed == "" {
			continue
		}
		result = append(result, trimmed)
	}
	return result
}

func appendIfMissing(items []string, value string) []string {
	for _, item := range items {
		if item == value {
			return items
		}
	}
	return append(items, value)
}

func containsFold(items []string, value string) bool {
	for _, item := range items {
		if strings.EqualFold(item, value) {
			return true
		}
	}
	return false
}

func allowedContains(allowed map[string]struct{}, value string) bool {
	if allowed == nil {
		return true
	}
	_, ok := allowed[strings.ToLower(strings.TrimSpace(value))]
	return ok
}