package providers

import (
	"regexp"
	"strings"
)

// ModelCapabilities describes what a model can do and its limits.
// Zero values mean "unknown / use heuristic".
type ModelCapabilities struct {
	Streaming      bool `json:"streaming"`
	ToolCalling    bool `json:"tool_calling"`
	Vision         bool `json:"vision"`
	Thinking       bool `json:"thinking"`
	Reasoning      bool `json:"reasoning"`
	CodeOptimized  bool `json:"code_optimized"`
	WebSearch      bool `json:"web_search"`
	MaxContext      int  `json:"max_context"`
	MaxOutputTokens int  `json:"max_output_tokens"`
}

// ── Known Model Defaults ────────────────────
// Authoritative output limits for models where the provider API
// doesn't report them. Keyed by exact model ID or prefix.
//
// Sources:
//   Anthropic: https://docs.anthropic.com/en/docs/about-claude/models
//   OpenAI:    https://platform.openai.com/docs/models
//   Meta:      Model cards on Hugging Face
//   Google:    https://ai.google.dev/gemini-api/docs/models

var knownModels = map[string]ModelCapabilities{
	// ── Anthropic ────────────────────────────
	"claude-opus-4": {
		Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
		MaxContext: 200000, MaxOutputTokens: 32000,
	},
	"claude-sonnet-4": {
		Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
		MaxContext: 200000, MaxOutputTokens: 16000,
	},
	"claude-3-5-sonnet": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 200000, MaxOutputTokens: 8192,
	},
	"claude-3-5-haiku": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 200000, MaxOutputTokens: 8192,
	},
	"claude-3-opus": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 200000, MaxOutputTokens: 4096,
	},
	"claude-3-haiku": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 200000, MaxOutputTokens: 4096,
	},

	// ── OpenAI ──────────────────────────────
	"gpt-4o": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 128000, MaxOutputTokens: 16384,
	},
	"gpt-4o-mini": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 128000, MaxOutputTokens: 16384,
	},
	"gpt-4-turbo": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 128000, MaxOutputTokens: 4096,
	},
	"gpt-4": {
		Streaming: true, ToolCalling: true,
		MaxContext: 8192, MaxOutputTokens: 4096,
	},
	"o1": {
		Streaming: true, Reasoning: true,
		MaxContext: 200000, MaxOutputTokens: 100000,
	},
	"o1-mini": {
		Streaming: true, Reasoning: true,
		MaxContext: 128000, MaxOutputTokens: 65536,
	},
	"o3": {
		Streaming: true, ToolCalling: true, Reasoning: true,
		MaxContext: 200000, MaxOutputTokens: 100000,
	},
	"o3-mini": {
		Streaming: true, ToolCalling: true, Reasoning: true,
		MaxContext: 200000, MaxOutputTokens: 65536,
	},
	"o4-mini": {
		Streaming: true, ToolCalling: true, Reasoning: true,
		MaxContext: 200000, MaxOutputTokens: 100000,
	},

	// ── Meta Llama ──────────────────────────
	"llama-3.1-405b": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 4096,
	},
	"llama-3.1-70b": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 4096,
	},
	"llama-3.1-8b": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 4096,
	},
	"llama-3.3-70b": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 4096,
	},
	"llama-4-maverick": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 1048576, MaxOutputTokens: 16384,
	},
	"llama-4-scout": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 524288, MaxOutputTokens: 16384,
	},

	// ── Google Gemini ───────────────────────
	"gemini-2.5-pro": {
		Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
		MaxContext: 1048576, MaxOutputTokens: 65536,
	},
	"gemini-2.5-flash": {
		Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
		MaxContext: 1048576, MaxOutputTokens: 65536,
	},
	"gemini-2.0-flash": {
		Streaming: true, ToolCalling: true, Vision: true,
		MaxContext: 1048576, MaxOutputTokens: 8192,
	},

	// ── DeepSeek ────────────────────────────
	"deepseek-r1": {
		Streaming: true, Reasoning: true,
		MaxContext: 65536, MaxOutputTokens: 8192,
	},
	"deepseek-v3": {
		Streaming: true, ToolCalling: true,
		MaxContext: 65536, MaxOutputTokens: 8192,
	},
	"deepseek-chat": {
		Streaming: true, ToolCalling: true,
		MaxContext: 65536, MaxOutputTokens: 8192,
	},

	// ── Mistral ─────────────────────────────
	"mistral-large": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 8192,
	},
	"mistral-small": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 8192,
	},
	"codestral": {
		Streaming: true, ToolCalling: true, CodeOptimized: true,
		MaxContext: 262144, MaxOutputTokens: 8192,
	},

	// ── Qwen ────────────────────────────────
	"qwen-2.5-72b": {
		Streaming: true, ToolCalling: true,
		MaxContext: 131072, MaxOutputTokens: 8192,
	},
	"qwen-2.5-coder-32b": {
		Streaming: true, ToolCalling: true, CodeOptimized: true,
		MaxContext: 131072, MaxOutputTokens: 8192,
	},
	"qwq-32b": {
		Streaming: true, Reasoning: true,
		MaxContext: 131072, MaxOutputTokens: 8192,
	},
}

// LookupKnownModel finds capabilities for a model by exact ID or prefix match.
// Returns the caps and true if found, zero value and false if not.
func LookupKnownModel(modelID string) (ModelCapabilities, bool) {
	id := strings.ToLower(modelID)

	// Strip provider prefix (e.g. "anthropic/claude-sonnet-4-20250514" → "claude-sonnet-4-20250514")
	if idx := strings.Index(id, "/"); idx >= 0 {
		id = id[idx+1:]
	}

	// Exact match first
	if caps, ok := knownModels[id]; ok {
		return caps, true
	}

	// Prefix match: "claude-sonnet-4-20250514" matches "claude-sonnet-4"
	var bestKey string
	var bestCaps ModelCapabilities
	for key, caps := range knownModels {
		if strings.HasPrefix(id, key) && len(key) > len(bestKey) {
			bestKey = key
			bestCaps = caps
		}
	}
	if bestKey != "" {
		return bestCaps, true
	}

	return ModelCapabilities{}, false
}

// ── Heuristic Capability Detection ──────────
// Ported from ai-editor's ProviderRegistry.
// Used for Ollama, LM Studio, and other generic endpoints
// that don't report capabilities in their /models response.

var (
	toolPatterns = []*regexp.Regexp{
		regexp.MustCompile(`(?i)llama[_-]?[34]`),
		regexp.MustCompile(`(?i)granite[_-]?[34]`),
		regexp.MustCompile(`(?i)hermes[_-]?[23]?`),
		regexp.MustCompile(`(?i)qwen[_-]?2`),
		regexp.MustCompile(`(?i)mistral|mixtral`),
		regexp.MustCompile(`(?i)command[_-]?r`),
		regexp.MustCompile(`(?i)phi[_-]?[34]`),
		regexp.MustCompile(`(?i)deepseek[_-]?v[23]`),
		regexp.MustCompile(`(?i)functionary`),
		regexp.MustCompile(`(?i)^gpt[_-]?[34]`),
		regexp.MustCompile(`(?i)^claude`),
		regexp.MustCompile(`(?i)gemma[_-]?2`),
		regexp.MustCompile(`(?i)glm[_-]?4`),
		regexp.MustCompile(`(?i)gemini`),
	}

	visionPatterns = []*regexp.Regexp{
		regexp.MustCompile(`(?i)llava|bakllava`),
		regexp.MustCompile(`(?i)moondream`),
		regexp.MustCompile(`(?i)llama.*vision|vision.*llama`),
		regexp.MustCompile(`(?i)minicpm[_-]?v`),
		regexp.MustCompile(`(?i)^gpt[_-]?4[_-]?o|^gpt[_-]?4[_-]?turbo`),
		regexp.MustCompile(`(?i)^claude[_-]?3`),
		regexp.MustCompile(`(?i)gemini`),
		regexp.MustCompile(`(?i)qwen.*vl|vl.*qwen`),
		regexp.MustCompile(`(?i)phi[_-]?[34].*vision`),
		regexp.MustCompile(`(?i)internvl`),
		regexp.MustCompile(`(?i)glm[_-]?4v`),
	}

	reasoningPatterns = []*regexp.Regexp{
		regexp.MustCompile(`(?i)deepseek[_-]?r1`),
		regexp.MustCompile(`(?i)qwq`),
		regexp.MustCompile(`(?i)^o[1234][_-]|^o[1234]$`),
	}

	codePatterns = []*regexp.Regexp{
		regexp.MustCompile(`(?i)codellama|code[_-]?llama`),
		regexp.MustCompile(`(?i)deepseek[_-]?coder`),
		regexp.MustCompile(`(?i)starcoder`),
		regexp.MustCompile(`(?i)coder|codestral`),
		regexp.MustCompile(`(?i)wizardcoder`),
		regexp.MustCompile(`(?i)codegemma`),
	}
)

func matchesAny(id string, patterns []*regexp.Regexp) bool {
	for _, p := range patterns {
		if p.MatchString(id) {
			return true
		}
	}
	return false
}

// InferCapabilities guesses model capabilities from the model ID string.
// This is the fallback when neither the known model table nor the provider
// API provides capability data.
func InferCapabilities(modelID string) ModelCapabilities {
	id := strings.ToLower(modelID)
	// Strip provider prefix
	if idx := strings.Index(id, "/"); idx >= 0 {
		id = id[idx+1:]
	}

	return ModelCapabilities{
		Streaming:     true, // virtually everything streams
		ToolCalling:   matchesAny(id, toolPatterns),
		Vision:        matchesAny(id, visionPatterns),
		Reasoning:     matchesAny(id, reasoningPatterns),
		CodeOptimized: matchesAny(id, codePatterns),
	}
}

// ResolveMaxOutput returns the max output tokens for a model.
// Priority:
//   1. Explicit value in caps (from DB / provider API)
//   2. Known model table
//   3. Derive from context window (context/8, clamped 2048..16384)
//   4. 4096 as absolute last resort
//
// This is the ONE place the default lives. Nothing else in the
// codebase should hardcode a max_tokens value.
func ResolveMaxOutput(modelID string, caps ModelCapabilities) int {
	// 1. Already set (from model_configs DB or provider API)
	if caps.MaxOutputTokens > 0 {
		return caps.MaxOutputTokens
	}

	// 2. Known model table
	if known, ok := LookupKnownModel(modelID); ok && known.MaxOutputTokens > 0 {
		return known.MaxOutputTokens
	}

	// 3. Derive from context window
	if caps.MaxContext > 0 {
		derived := caps.MaxContext / 8
		if derived < 2048 {
			derived = 2048
		}
		if derived > 16384 {
			derived = 16384
		}
		return derived
	}

	// 4. Last resort — the ONLY place 4096 appears as a default
	return 4096
}

// MergeCapabilities takes authoritative caps (from provider API or DB) and fills
// gaps from the known model table and heuristic detection. Provider-reported data
// always wins; known table fills missing fields; heuristics are last resort.
func MergeCapabilities(authoritative ModelCapabilities, modelID string) ModelCapabilities {
	merged := authoritative

	// Fill gaps from known model table
	known, found := LookupKnownModel(modelID)
	if found {
		if !merged.ToolCalling && known.ToolCalling {
			merged.ToolCalling = true
		}
		if !merged.Vision && known.Vision {
			merged.Vision = true
		}
		if !merged.Thinking && known.Thinking {
			merged.Thinking = true
		}
		if !merged.Reasoning && known.Reasoning {
			merged.Reasoning = true
		}
		if !merged.CodeOptimized && known.CodeOptimized {
			merged.CodeOptimized = true
		}
		if !merged.WebSearch && known.WebSearch {
			merged.WebSearch = true
		}
		if merged.MaxContext == 0 && known.MaxContext > 0 {
			merged.MaxContext = known.MaxContext
		}
		if merged.MaxOutputTokens == 0 && known.MaxOutputTokens > 0 {
			merged.MaxOutputTokens = known.MaxOutputTokens
		}
		return merged
	}

	// No known model — fill gaps from heuristics
	inferred := InferCapabilities(modelID)
	if !merged.ToolCalling && inferred.ToolCalling {
		merged.ToolCalling = true
	}
	if !merged.Vision && inferred.Vision {
		merged.Vision = true
	}
	if !merged.Thinking && inferred.Thinking {
		merged.Thinking = true
	}
	if !merged.Reasoning && inferred.Reasoning {
		merged.Reasoning = true
	}
	if !merged.CodeOptimized && inferred.CodeOptimized {
		merged.CodeOptimized = true
	}
	if merged.MaxContext == 0 && inferred.MaxContext > 0 {
		merged.MaxContext = inferred.MaxContext
	}
	if merged.MaxOutputTokens == 0 && inferred.MaxOutputTokens > 0 {
		merged.MaxOutputTokens = inferred.MaxOutputTokens
	}

	return merged
}

// HasProviderData returns true if this capability set contains any data that
// was likely reported by a provider (not just zero values).
func (c ModelCapabilities) HasProviderData() bool {
	return c.ToolCalling || c.Vision || c.Thinking || c.Reasoning ||
		c.CodeOptimized || c.WebSearch || c.MaxContext > 0 || c.MaxOutputTokens > 0
}