package config

import (
	"fmt"
	"strings"

	"iop/packages/go/execution"
)

// Category represents the provider category.
type Category string

const (
	CategoryAPI            Category = "api"
	CategoryLocalInference Category = "local_inference"
)

// validateCategory returns nil when category is valid.
func validateCategory(c Category) error {
	switch c {
	case CategoryAPI, CategoryLocalInference:
		return nil
	default:
		return fmt.Errorf("invalid provider category %q (allowed: %q, %q)", c, CategoryAPI, CategoryLocalInference)
	}
}

// NodeProviderConf defines a resource/provider candidate that belongs to a
// node. Provider-first resources can declare type-specific execution fields
// directly, while Adapter remains available for legacy/compat overrides.
type NodeProviderConf struct {
	// ID uniquely identifies this provider within the node and is referenced
	// by models[].providers keys.
	ID string `mapstructure:"id" yaml:"id"`
	// Type is the runtime type (e.g. "ollama", "vllm", "vllm-mlx", "lemonade", "sglang", "openai_api").
	Type string `mapstructure:"type" yaml:"type"`
	// Category classifies the resource; values: api, local_inference.
	Category Category `mapstructure:"category" yaml:"category"`
	// Adapter is the optional concrete enabled adapter instance key used by
	// legacy/compat provider resources. A legacy type-name route is valid only
	// when exactly one enabled adapter instance of that type exists on the same
	// node.
	Adapter string `mapstructure:"adapter" yaml:"adapter,omitempty"`
	// Models lists the served model names this provider can actually serve.
	Models []string `mapstructure:"models" yaml:"models,omitempty"`
	// Health is the observed provider health state.
	Health string `mapstructure:"health" yaml:"health,omitempty"`
	// Capacity is the provider/resource maximum concurrent execution slots.
	// Provider-pool dispatch treats 0 as not dispatchable rather than falling
	// back to node runtime concurrency metadata.
	Capacity int `mapstructure:"capacity" yaml:"capacity,omitempty"`
	// TotalContextTokens is the provider runtime's total KV/context budget used
	// as the long-context admission budget. Runtime-specific options (e.g.
	// ctx_size) map onto this value. Must satisfy
	// total_context_tokens >= context_window_tokens * long_context_capacity for
	// every model group that references this provider.
	TotalContextTokens int `mapstructure:"total_context_tokens" yaml:"total_context_tokens,omitempty"`
	// LongContextCapacity is the number of concurrent long-context slots this
	// provider allows (context_window_tokens-sized requests). Long requests
	// occupy both a normal capacity slot and a long slot. Zero means the
	// provider declares no dedicated long-slot limit.
	LongContextCapacity int `mapstructure:"long_context_capacity" yaml:"long_context_capacity,omitempty"`
	// MaxQueue is a legacy per-provider queue depth field. The effective
	// policy owner is the Edge root provider_pool.max_queue.
	// This field is retained only for compatibility and is ignored after
	// LoadEdge migration once EdgeProviderPoolConf is resolved.
	MaxQueue int `mapstructure:"max_queue" yaml:"max_queue,omitempty"`
	// QueueTimeoutMS is a legacy per-provider queue timeout field. The
	// effective policy owner is the Edge root provider_pool.queue_timeout_ms.
	// This field is retained only for compatibility and is ignored after
	// LoadEdge migration once EdgeProviderPoolConf is resolved.
	QueueTimeoutMS int `mapstructure:"queue_timeout_ms" yaml:"queue_timeout_ms,omitempty"`
	// Priority is the provider dispatch priority. Lower values are preferred
	// when in_flight counts are equal. Must be non-negative; zero is the
	// default priority.
	Priority int `mapstructure:"priority" yaml:"priority,omitempty"`
	// LifecycleCapabilities lists coarse lifecycle capabilities.
	LifecycleCapabilities []string `mapstructure:"lifecycle_capabilities" yaml:"lifecycle_capabilities,omitempty"`
	// Enabled controls whether this provider participates in the dispatch pool.
	// Nil (omitted) means enabled; false disables the provider without removing it.
	Enabled *bool `mapstructure:"enabled" yaml:"enabled,omitempty"`

	// Profile is the optional protocol profile selector. When set, it selects
	// a built-in or custom profile from the top-level protocol_profiles catalog.
	// When empty, the provider type is used for legacy alias normalization.
	Profile string `mapstructure:"profile" yaml:"profile,omitempty"`
	// RuntimeProfile is the resolved concrete protocol profile snapshot.
	// It is populated during config normalization and is excluded from YAML
	// serialization (it is a runtime-only field). It is never written back
	// to source config.
	RuntimeProfile *ConcreteProtocolProfile `mapstructure:"-" yaml:"-"`

	// Provider-First fields (G06)
	Provider         string            `mapstructure:"provider" yaml:"provider,omitempty"`
	Endpoint         string            `mapstructure:"endpoint" yaml:"endpoint,omitempty"`
	BaseURL          string            `mapstructure:"base_url" yaml:"base_url,omitempty"`
	Headers          map[string]string `mapstructure:"headers" yaml:"headers,omitempty"`
	ContextSize      int               `mapstructure:"context_size" yaml:"context_size,omitempty"`
	RequestTimeoutMS int               `mapstructure:"request_timeout_ms" yaml:"request_timeout_ms,omitempty"`
	// ResponseStallTimeoutMS is the provider-originated response-stall timeout
	// in milliseconds. It is carried on every provider-first and legacy route
	// request so the downstream watchdog has one effective value per dispatched
	// attempt. Zero is treated as the documented default (60000 ms). Negative
	// values and positive values that cannot safely become a time.Duration are
	// rejected by Validate. The effective value is returned by
	// EffectiveResponseStallTimeoutMS.
	ResponseStallTimeoutMS int64 `mapstructure:"response_stall_timeout_ms" yaml:"response_stall_timeout_ms,omitempty"`
}

// EffectiveResponseStallTimeoutMS returns the effective response-stall timeout
// in milliseconds. Validate rejects negative and overflow values at load, so
// here only zero maps to the shared default and safe positives pass through.
func (p NodeProviderConf) EffectiveResponseStallTimeoutMS() int64 {
	if p.ResponseStallTimeoutMS == 0 {
		return execution.DefaultResponseStallTimeoutMS
	}
	return p.ResponseStallTimeoutMS
}

// Validate checks internal consistency of the provider candidate config.
func (p NodeProviderConf) Validate() error {
	id := strings.TrimSpace(p.ID)
	if id == "" {
		return fmt.Errorf("nodes[].providers[].id must not be empty")
	}
	if strings.TrimSpace(p.Type) == "" {
		return fmt.Errorf("nodes[].providers[%q]: type must not be empty", id)
	}
	if err := validateCategory(p.Category); err != nil {
		return fmt.Errorf("nodes[].providers[%q]: %w", id, err)
	}
	for i, m := range p.Models {
		if strings.TrimSpace(m) == "" {
			return fmt.Errorf("nodes[].providers[%q].models[%d]: served model name must not be empty", id, i)
		}
	}
	if p.Capacity < 0 {
		return fmt.Errorf("nodes[].providers[%q].capacity must be non-negative", id)
	}
	if p.MaxQueue < 0 {
		return fmt.Errorf("nodes[].providers[%q].max_queue must be non-negative", id)
	}
	if p.QueueTimeoutMS < 0 {
		return fmt.Errorf("nodes[].providers[%q].queue_timeout_ms must be non-negative", id)
	}
	if p.RequestTimeoutMS < 0 {
		return fmt.Errorf("nodes[].providers[%q].request_timeout_ms must be non-negative", id)
	}
	if p.ContextSize < 0 {
		return fmt.Errorf("nodes[].providers[%q].context_size must be non-negative", id)
	}
	if p.Priority < 0 {
		return fmt.Errorf("nodes[].providers[%q].priority must be non-negative", id)
	}
	if p.TotalContextTokens < 0 {
		return fmt.Errorf("nodes[].providers[%q].total_context_tokens must be non-negative", id)
	}
	if p.LongContextCapacity < 0 {
		return fmt.Errorf("nodes[].providers[%q].long_context_capacity must be non-negative", id)
	}
	if p.LongContextCapacity > 0 && p.TotalContextTokens <= 0 {
		return fmt.Errorf("nodes[].providers[%q].total_context_tokens must be positive when long_context_capacity > 0", id)
	}
	if err := validateStallTimeout(p.ResponseStallTimeoutMS, id); err != nil {
		return err
	}

	return nil
}

// validateStallTimeout enforces the response_stall_timeout_ms contract:
// zero selects the default, positive values must fit safely into a
// time.Duration, and negative values are rejected.
func validateStallTimeout(ms int64, id string) error {
	if err := execution.ValidateStallTimeoutMS(ms); err != nil {
		return fmt.Errorf("nodes[].providers[%q].response_stall_timeout_ms: %w", id, err)
	}
	return nil
}

// ProviderEnabled reports whether a provider is enabled.
// Returns true when Enabled is nil (omitted) or *Enabled is true.
// Returns false only when Enabled is explicitly set to false.
func ProviderEnabled(p NodeProviderConf) bool {
	return p.Enabled == nil || *p.Enabled
}

const (
	UsageAttributionProvider   = "provider"
	UsageAttributionModelGroup = "model_group"
)

// ModelCatalogEntry is a top-level Edge config entry that defines a canonical
// routing key (`ID`) and its provider-pool mapping.  Each provider id key
// maps to the concrete served model name that the provider actually exposes.
// Exactly one of Providers or ExecutionPreset must be set: provider-only
// entries continue to dispatch through the provider pool, while preset-only
// entries (virtual models) bind to a single frozen execution preset shape.
type ModelCatalogEntry struct {
	// ID is the canonical routing key (e.g. "qwen3.6:35b") that matches the
	// external OpenAI-compatible model field.
	ID string `mapstructure:"id" yaml:"id"`
	// UsageAttribution selects the approved usage aggregation basis. Omission
	// defaults to the actual provider; model_group is an explicit opt-in for
	// operators that have approved every candidate as the same logical model.
	UsageAttribution string `mapstructure:"usage_attribution" yaml:"usage_attribution,omitempty"`
	// DisplayName is a human-readable name shown in operation dashboards.
	DisplayName string `mapstructure:"display_name" yaml:"display_name,omitempty"`
	// ContextWindowTokens is the model group's single-request maximum context
	// contract shared by every provider in the group.
	ContextWindowTokens int `mapstructure:"context_window_tokens" yaml:"context_window_tokens,omitempty"`
	// DefaultMaxTokens is injected when OpenAI-compatible requests omit an
	// output-token limit for this model group.
	DefaultMaxTokens int `mapstructure:"default_max_tokens" yaml:"default_max_tokens,omitempty"`
	// MinMaxTokens raises too-small OpenAI-compatible output-token limits for
	// this model group.
	MinMaxTokens int `mapstructure:"min_max_tokens" yaml:"min_max_tokens,omitempty"`
	// DefaultThinkingTokenBudget bounds reasoning tokens for think-capable
	// providers when callers do not explicitly set a budget.
	DefaultThinkingTokenBudget int `mapstructure:"default_thinking_token_budget" yaml:"default_thinking_token_budget,omitempty"`
	// Providers maps provider id to the concrete served model name that the
	// provider actually exposes.  Keys must match nodes[].providers[].id.
	Providers map[string]string `mapstructure:"providers" yaml:"providers"`
	// ExecutionPreset is the stable execution preset id this model binds to.
	// When set, the model is a virtual entry that delegates execution to the
	// named preset; Providers must be empty and provider-only budget/token
	// checks are skipped for the entry.
	ExecutionPreset string `mapstructure:"execution_preset" yaml:"execution_preset,omitempty"`
	// TokenCounter declares how a model group's input tokens are counted
	// without an upstream call. Only valid for Chat-only profiles.
	TokenCounter *TokenCounterConf `mapstructure:"token_counter" yaml:"token_counter,omitempty"`
}

// EffectiveUsageAttribution returns the validated attribution policy, applying
// the provider default when the field is omitted.
func (e ModelCatalogEntry) EffectiveUsageAttribution() string {
	if attribution := strings.TrimSpace(e.UsageAttribution); attribution != "" {
		return attribution
	}
	return UsageAttributionProvider
}

// buildProviderServedModelsIndex aggregates all nodes[].providers[].models
// keyed by provider id. Returns the index for Validate.
func buildProviderServedModelsIndex(nodes []NodeDefinition) map[string]map[string]struct{} {
	serveModels := make(map[string]map[string]struct{})
	for i := range nodes {
		for _, p := range nodes[i].Providers {
			if _, ok := serveModels[p.ID]; !ok {
				serveModels[p.ID] = make(map[string]struct{})
			}
			for _, m := range p.Models {
				serveModels[p.ID][m] = struct{}{}
			}
		}
	}
	return serveModels
}

// Validate checks the structure and internal references of the model catalog entry.
// When serveModels is non-nil, it also verifies that each models[].providers[pid]=model
// refers to a served model that exists in nodes[].providers[id].models.
func (e ModelCatalogEntry) Validate(resolvedProviderIDs map[string]struct{}, serveModels map[string]map[string]struct{}) error {
	id := strings.TrimSpace(e.ID)
	if id == "" {
		return fmt.Errorf("models[].id must not be empty")
	}
	switch attribution := e.EffectiveUsageAttribution(); attribution {
	case UsageAttributionProvider, UsageAttributionModelGroup:
	default:
		return fmt.Errorf("models[%q].usage_attribution must be %q or %q, got %q",
			id, UsageAttributionProvider, UsageAttributionModelGroup, attribution)
	}
	if e.DefaultMaxTokens < 0 {
		return fmt.Errorf("models[%q].default_max_tokens must be non-negative", id)
	}
	if e.MinMaxTokens < 0 {
		return fmt.Errorf("models[%q].min_max_tokens must be non-negative", id)
	}
	if e.DefaultThinkingTokenBudget < 0 {
		return fmt.Errorf("models[%q].default_thinking_token_budget must be non-negative", id)
	}
	if e.ContextWindowTokens < 0 {
		return fmt.Errorf("models[%q].context_window_tokens must be non-negative", id)
	}
	if e.DefaultMaxTokens > 0 && e.MinMaxTokens > 0 && e.DefaultMaxTokens < e.MinMaxTokens {
		return fmt.Errorf("models[%q].default_max_tokens must be greater than or equal to min_max_tokens", id)
	}
	// Enforce exactly one of Providers or ExecutionPreset.
	isVirtual := strings.TrimSpace(e.ExecutionPreset) != ""
	if len(e.Providers) == 0 && !isVirtual {
		return fmt.Errorf("models[%q].providers must not be empty", id)
	}
	if len(e.Providers) > 0 && isVirtual {
		return fmt.Errorf("models[%q]: exactly one of providers or execution_preset must be set, got both", id)
	}
	if isVirtual {
		return nil // preset-only virtual entry; preset id resolved later by LoadEdge
	}
	for pid, model := range e.Providers {
		p := strings.TrimSpace(pid)
		if p == "" {
			return fmt.Errorf("models[%q].providers: provider id must not be empty", id)
		}
		m := strings.TrimSpace(model)
		if m == "" {
			return fmt.Errorf("models[%q].providers[%q]: served model name must not be empty", id, p)
		}
		if _, ok := resolvedProviderIDs[p]; !ok {
			return fmt.Errorf("models[%q].providers[%q]: provider id %q does not resolve to any nodes[].providers[].id", id, p, p)
		}
		// Serve model membership validation (REVIEW_API-1):
		// The served model must be listed in the provider's own models[].
		if serveModels != nil {
			servedSet, exists := serveModels[p]
			if !exists {
				// No providers list yet (should not happen in modern configs).
				return fmt.Errorf("models[%q].providers[%q]: provider %q has no models[] defined", id, p, p)
			}
			if _, ok := servedSet[m]; !ok {
				return fmt.Errorf("models[%q].providers[%q]: served model %q is not listed in nodes[].providers[%q].models", id, p, m, p)
			}
		}
	}
	return nil
}
