/*
Copyright 2026 LocoStack.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package litellm

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/LocoStack/loco-operator/api/v1alpha1"
)

type LiteLLMConfig struct {
	ModelList       []modelEntry              `yaml:"model_list"`
	MCPServers      map[string]mcpServerEntry `yaml:"mcp_servers,omitempty"`
	LiteLLMSettings *liteLLMSettings          `yaml:"litellm_settings,omitempty"`
	GeneralSettings *liteLLMGeneralSettings   `yaml:"general_settings,omitempty"`
	EnvVars         map[string]string         `yaml:"environment_variables,omitempty"`
}

type modelEntry struct {
	ModelName     string         `yaml:"model_name"`
	LiteLLMParams map[string]any `yaml:"litellm_params"`
}

type mcpServerEntry struct {
	// url is the HTTP/SSE endpoint for sse and http transports.
	URL string `yaml:"url,omitempty"`
	// transport selects the MCP transport: sse (default), http, or stdio.
	Transport string `yaml:"transport,omitempty"`

	// --- authentication ---

	// authType selects the auth method: api_key, bearer_token, basic,
	// authorization, oauth2, or aws_sigv4.
	AuthType string `yaml:"auth_type,omitempty"`
	// authValue is the credential for api_key, bearer_token, basic, and
	// authorization auth types. May reference an env var via os.environ/VAR.
	AuthValue string `yaml:"auth_value,omitempty"`

	// --- headers ---

	// staticHeaders are key-value pairs sent with every request to this server.
	StaticHeaders map[string]string `yaml:"static_headers,omitempty"`
	// extraHeaders lists client request header names that LiteLLM should
	// forward to this MCP server.
	ExtraHeaders []string `yaml:"extra_headers,omitempty"`

	// allowAllKeys grants every LiteLLM API key access to this server when true.
	AllowAllKeys bool `yaml:"allow_all_keys,omitempty"`
}

type liteLLMSettings struct {
	SuccessCallback []string `yaml:"success_callback,omitempty"`
	FailureCallback []string `yaml:"failure_callback,omitempty"`
	JSONLogs        bool     `yaml:"json_logs,omitempty"`
}

type liteLLMGeneralSettings struct {
	MasterKey string `yaml:"master_key,omitempty"`
}

type LiteLLMConfigBuilder struct {
	MasterKeyEnvName       string
	Stack                  *v1alpha1.Stack
	ObservabilityComponent *v1alpha1.Component
	ExternalModels         []*v1alpha1.ExternalModel
	ManagedModels          []*v1alpha1.ManagedModel
	ExternalTools          []*v1alpha1.ExternalTool
	ManagedTools           []*v1alpha1.ManagedTool
}

func (l *LiteLLMConfigBuilder) BuildLiteLLMConfig() (LiteLLMConfig, error) {
	modelList := []modelEntry{}
	for _, em := range l.ExternalModels {
		entry := buildExternalModelEntry(em)
		modelList = append(modelList, entry)
	}
	for _, mm := range l.ManagedModels {
		if mm.Status.Endpoint == "" {
			continue
		}
		entry := buildManagedModelEntry(mm)
		modelList = append(modelList, entry)
	}
	mcpServers := map[string]mcpServerEntry{}
	for _, et := range l.ExternalTools {
		entry := buildExternalToolEntry(et)
		mcpServers[et.Spec.ToolName] = entry
	}
	for _, mt := range l.ManagedTools {
		if mt.Status.Endpoint == "" {
			continue
		}
		entry := buildManagedToolEntry(mt)
		mcpServers[mt.Spec.ToolName] = entry
	}
	phoenixEnabled := phoenixEnabled(l.Stack)
	successCB := []string{}
	failureCB := []string{}
	if phoenixEnabled {
		successCB = append(successCB, "arize_phoenix")
		failureCB = append(failureCB, "arize_phoenix")
	}
	cfg := LiteLLMConfig{
		ModelList:  modelList,
		MCPServers: mcpServers,
		LiteLLMSettings: &liteLLMSettings{
			SuccessCallback: successCB,
			FailureCallback: failureCB,
			JSONLogs:        true,
		},
		GeneralSettings: &liteLLMGeneralSettings{
			MasterKey: "os.environ/" + l.MasterKeyEnvName,
		},
	}
	if phoenixEnabled {
		cfg.EnvVars = map[string]string{
			"PHOENIX_PROJECT_NAME":            "locostack",
			"PHOENIX_COLLECTOR_HTTP_ENDPOINT": l.ObservabilityComponent.Status.Endpoint + "/v1/traces",
		}
	}
	return cfg, nil
}

func phoenixEnabled(stack *v1alpha1.Stack) bool {
	if stack == nil {
		return false
	}
	if stack.Spec.Observability == nil {
		return false
	}
	if !stack.Spec.Observability.Enabled {
		return false
	}
	if stack.Spec.Observability.Template != nil && stack.Spec.Observability.Template.Name != "phoenix" {
		return false
	}
	return true
}

func AuthEnvVarName(kind string, name string) string {
	return fmt.Sprintf("%s_%s_AUTH", strings.ToUpper(kind), strings.ToUpper(strings.ReplaceAll(name, "-", "_")))
}

// buildExternalModelEntry converts an ExternalModel into a LiteLLM model_list entry.
func buildExternalModelEntry(em *v1alpha1.ExternalModel) modelEntry {
	params := map[string]any{
		"model": fmt.Sprintf("%s/%s", em.Spec.Provider, em.Spec.ProviderModel),
	}
	if em.Spec.Auth != nil && (em.Spec.Auth.APIKey != nil || em.Spec.Auth.BearerToken != nil) {
		params["api_key"] = "os.environ/" + AuthEnvVarName("ExternalModel", em.Name)
	}
	if em.Spec.APIBase != "" {
		params["api_base"] = em.Spec.APIBase
	}
	if em.Spec.APIVersion != "" {
		params["api_version"] = em.Spec.APIVersion
	}
	switch em.Spec.Category {
	case "embedding":
		params["mode"] = "embedding"
	case "reranker":
		params["mode"] = "rerank"
	}
	for k, v := range em.Spec.ExtraParams {
		params[k] = v
	}
	if em.Spec.DefaultInferenceParams != nil {
		applyInferenceParams(params, em.Spec.DefaultInferenceParams)
	}
	return modelEntry{
		ModelName:     em.Spec.ModelName,
		LiteLLMParams: params,
	}
}

// buildManagedModelEntry converts a ManagedModel into a LiteLLM model_list entry.
func buildManagedModelEntry(mm *v1alpha1.ManagedModel) modelEntry {
	// litellm's rerank dispatch does not recognize "openai" as a provider;
	// "hosted_vllm" is the generic self-hosted provider for the /rerank call type.
	provider := "openai"
	if mm.Spec.Category == "reranker" {
		provider = "hosted_vllm"
	}
	params := map[string]any{
		"model":    fmt.Sprintf("%s/%s", provider, mm.Spec.ModelName),
		"api_base": mm.Status.Endpoint + "/v1",
		"api_key":  "not-set",
	}
	switch mm.Spec.Category {
	case "embedding":
		params["mode"] = "embedding"
	case "reranker":
		params["mode"] = "rerank"
	}
	return modelEntry{
		ModelName:     mm.Spec.ModelName,
		LiteLLMParams: params,
	}
}

func applyInferenceParams(params map[string]any, inf *v1alpha1.InferenceParameters) {
	if inf.Temperature != "" {
		if v, err := strconv.ParseFloat(inf.Temperature, 64); err == nil {
			params["temperature"] = v
		}
	}
	if inf.TopP != "" {
		if v, err := strconv.ParseFloat(inf.TopP, 64); err == nil {
			params["top_p"] = v
		}
	}
	if inf.MaxTokens != nil {
		params["max_tokens"] = *inf.MaxTokens
	}
	for k, v := range inf.Extra {
		params[k] = v
	}
}

// buildExternalToolEntry converts an ExternalTool into a LiteLLM mcp_server entry.
func buildExternalToolEntry(et *v1alpha1.ExternalTool) mcpServerEntry {
	return buildToolEntry("ExternalTool", et.Spec.ToolName, et.Spec.Endpoint, et.Spec.Transport, et.Spec.Auth)
}

// buildManagedToolEntry converts a ManagedTool into a LiteLLM mcp_server entry.
func buildManagedToolEntry(mt *v1alpha1.ManagedTool) mcpServerEntry {
	url := mt.Status.Endpoint
	switch mt.Spec.Transport {
	case "http":
		url = fmt.Sprintf("%s/mcp", strings.TrimSuffix(mt.Status.Endpoint, "/"))
	case "sse":
		url = fmt.Sprintf("%s/sse", strings.TrimSuffix(mt.Status.Endpoint, "/"))
	}
	return buildToolEntry("ManagedTool", mt.Spec.ToolName, url, mt.Spec.Transport, mt.Spec.Auth)
}

func buildToolEntry(toolKind string, toolName string, endpoint string, transport string, auth *v1alpha1.AuthSpec) mcpServerEntry {
	entry := mcpServerEntry{
		URL:       endpoint,
		Transport: transport,
	}
	if auth == nil {
		return entry
	}
	if auth.BearerToken != nil {
		entry.AuthType = "bearer_token"
		entry.AuthValue = "os.environ/" + AuthEnvVarName(toolKind, toolName)
	} else if auth.APIKey != nil {
		headerName := auth.APIKey.HeaderName
		if headerName == "" {
			headerName = "X-Api-Key"
		}
		if entry.StaticHeaders == nil {
			entry.StaticHeaders = make(map[string]string)
		}
		entry.StaticHeaders[headerName] = "os.environ/" + AuthEnvVarName(toolKind, toolName)
	}

	for _, h := range auth.Headers {
		if entry.StaticHeaders == nil {
			entry.StaticHeaders = make(map[string]string)
		}
		entry.StaticHeaders[h.Name] = "os.environ/" + AuthEnvVarName(toolKind, toolName+"_"+h.Name)
	}

	return entry
}