package capabilities

import (
	"regexp"
	"strings"

	"git.gobha.me/xcaliber/chat-switchboard/models"
)

// ── Known Model Defaults ────────────────────
// Authoritative capabilities 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]models.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.
func LookupKnownModel(modelID string) (models.ModelCapabilities, bool) {
	id := normalizeModelID(modelID)

	// 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 models.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 models.ModelCapabilities{}, false
}

// ── Heuristic Capability Detection ──────────

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.
// Fallback when neither the known model table nor the provider API has data.
func InferCapabilities(modelID string) models.ModelCapabilities {
	id := normalizeModelID(modelID)
	return models.ModelCapabilities{
		Streaming:     true, // virtually everything streams
		ToolCalling:   matchesAny(id, toolPatterns),
		Vision:        matchesAny(id, visionPatterns),
		Reasoning:     matchesAny(id, reasoningPatterns),
		CodeOptimized: matchesAny(id, codePatterns),
	}
}

// ResolveIntrinsic determines the intrinsic capabilities of a model.
// Priority:
//  1. catalogCaps (from model_catalog DB — provider API data)
//  2. Known model table (static, compiled-in)
//  3. Heuristic inference (regex patterns)
//
// This is the ONLY function that computes intrinsic capabilities.
func ResolveIntrinsic(modelID string, catalogCaps *models.ModelCapabilities) models.ModelCapabilities {
	// Start with catalog data if available
	var base models.ModelCapabilities
	if catalogCaps != nil && catalogCaps.HasProviderData() {
		base = *catalogCaps
	}

	// Fill gaps from known model table
	if known, found := LookupKnownModel(modelID); found {
		mergeGaps(&base, &known)
		return base
	}

	// Fill gaps from heuristic inference
	inferred := InferCapabilities(modelID)
	mergeGaps(&base, &inferred)
	return base
}

// ResolveMaxOutput returns the max output tokens for a model.
// Priority: explicit caps → known model table → derive from context → 4096 fallback.
func ResolveMaxOutput(modelID string, caps models.ModelCapabilities) int {
	if caps.MaxOutputTokens > 0 {
		return caps.MaxOutputTokens
	}
	if known, ok := LookupKnownModel(modelID); ok && known.MaxOutputTokens > 0 {
		return known.MaxOutputTokens
	}
	if caps.MaxContext > 0 {
		derived := caps.MaxContext / 8
		if derived < 2048 {
			derived = 2048
		}
		if derived > 16384 {
			derived = 16384
		}
		return derived
	}
	return 4096
}

// mergeGaps fills zero/false fields in dst from src. Never overrides existing data.
func mergeGaps(dst, src *models.ModelCapabilities) {
	if !dst.Streaming && src.Streaming {
		dst.Streaming = true
	}
	if !dst.ToolCalling && src.ToolCalling {
		dst.ToolCalling = true
	}
	if !dst.Vision && src.Vision {
		dst.Vision = true
	}
	if !dst.Thinking && src.Thinking {
		dst.Thinking = true
	}
	if !dst.Reasoning && src.Reasoning {
		dst.Reasoning = true
	}
	if !dst.CodeOptimized && src.CodeOptimized {
		dst.CodeOptimized = true
	}
	if !dst.WebSearch && src.WebSearch {
		dst.WebSearch = true
	}
	if dst.MaxContext == 0 && src.MaxContext > 0 {
		dst.MaxContext = src.MaxContext
	}
	if dst.MaxOutputTokens == 0 && src.MaxOutputTokens > 0 {
		dst.MaxOutputTokens = src.MaxOutputTokens
	}
}

// normalizeModelID strips provider prefix and lowercases.
func normalizeModelID(modelID string) string {
	id := strings.ToLower(modelID)
	if idx := strings.Index(id, "/"); idx >= 0 {
		id = id[idx+1:]
	}
	return id
}