package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const (
	modelsDevURL             = "https://models.dev/api.json"
	novitaModelsURL          = "https://api.novita.ai/openai/v1/models"
	configMapName            = "agentgateway-model-costs"
	serviceAccountDir        = "/var/run/secrets/kubernetes.io/serviceaccount"
	defaultSyncInterval      = "6h"
	defaultSyncRetryInterval = "5m"
	defaultEnabledProviders  = "google-vertex,novita"
)

type modelsDevCatalog map[string]modelsDevProvider

type modelsDevProvider struct {
	Models map[string]modelsDevModel `json:"models"`
}

type modelsDevModel struct {
	Cost modelsDevCost `json:"cost"`
}

type modelsDevCost struct {
	Input       *json.Number    `json:"input"`
	Output      *json.Number    `json:"output"`
	CacheRead   *json.Number    `json:"cache_read"`
	CacheWrite  *json.Number    `json:"cache_write"`
	Reasoning   *json.Number    `json:"reasoning"`
	InputAudio  *json.Number    `json:"input_audio"`
	OutputAudio *json.Number    `json:"output_audio"`
	Tiers       []modelsDevTier `json:"tiers"`
}

type modelsDevTier struct {
	Input       *json.Number          `json:"input"`
	Output      *json.Number          `json:"output"`
	CacheRead   *json.Number          `json:"cache_read"`
	CacheWrite  *json.Number          `json:"cache_write"`
	Reasoning   *json.Number          `json:"reasoning"`
	InputAudio  *json.Number          `json:"input_audio"`
	OutputAudio *json.Number          `json:"output_audio"`
	Tier        modelsDevTierSelector `json:"tier"`
}

type modelsDevTierSelector struct {
	Type string `json:"type"`
	Size int    `json:"size"`
}

type novitaCatalog struct {
	Data []novitaModel `json:"data"`
}

type novitaModel struct {
	ID          string      `json:"id"`
	InputPrice  json.Number `json:"input_token_price_per_m"`
	OutputPrice json.Number `json:"output_token_price_per_m"`
}

type catalog struct {
	Providers map[string]catalogProvider `json:"providers"`
}

type catalogProvider struct {
	Models map[string]catalogModel `json:"models"`
}

type catalogModel struct {
	Rates map[string]string `json:"rates"`
	Tiers []catalogTier     `json:"tiers,omitempty"`
}

type catalogTier struct {
	ContextOver int               `json:"contextOver"`
	Rates       map[string]string `json:"rates"`
}

type configMapPatch struct {
	Data map[string]string `json:"data"`
}

type configMapResource struct {
	APIVersion string            `json:"apiVersion"`
	Kind       string            `json:"kind"`
	Metadata   configMapMetadata `json:"metadata"`
	Data       map[string]string `json:"data"`
}

type configMapMetadata struct {
	Name        string            `json:"name"`
	Annotations map[string]string `json:"annotations"`
	Labels      map[string]string `json:"labels"`
}

type modelList struct {
	Object string          `json:"object"`
	Data   []modelListItem `json:"data"`
}

type modelListItem struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

type providerDefinition struct {
	CatalogID      string
	ModelsDevID    string
	ModelPrefix    string
	SourceProvider string
	Novita         bool
}

type providerSelection map[string]providerDefinition

var modelsDevProviderIDs = map[string]string{
	"amazon-bedrock": "aws.bedrock",
	"anthropic":      "anthropic",
	"azure":          "azure",
	"baseten":        "baseten",
	"cerebras":       "cerebras",
	"cohere":         "cohere",
	"deepinfra":      "deepinfra",
	"deepseek":       "deepseek",
	"fireworks-ai":   "fireworks",
	"github-copilot": "copilot",
	"google":         "gcp.gemini",
	"google-vertex":  "gcp.vertex_ai",
	"groq":           "groq",
	"huggingface":    "huggingface",
	"mistral":        "mistral",
	"openai":         "openai",
	"openrouter":     "openrouter",
	"togetherai":     "togetherai",
	"xai":            "xai",
}

func main() {
	var err error
	if len(os.Args) != 2 {
		err = errors.New("expected sync or run")
	} else {
		var providers providerSelection
		providers, err = providersFromEnvironment()
		if err == nil {
			switch os.Args[1] {
			case "sync":
				err = syncCatalog(context.Background(), providers)
			case "run":
				err = run(context.Background(), providers)
			default:
				err = fmt.Errorf("unknown command %q", os.Args[1])
			}
		}
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func providersFromEnvironment() (providerSelection, error) {
	value := os.Getenv("ENABLED_PROVIDERS")
	if value == "" {
		value = defaultEnabledProviders
	}
	definitions := buildProviderDefinitions()
	providers := providerSelection{}
	for _, name := range strings.Split(value, ",") {
		name = strings.TrimSpace(name)
		definition, found := definitions[name]
		if !found {
			return nil, fmt.Errorf("unsupported provider %q", name)
		}
		if _, found := providers[name]; found {
			return nil, fmt.Errorf("duplicate provider %q", name)
		}
		providers[name] = definition
	}
	if len(providers) == 0 {
		return nil, errors.New("ENABLED_PROVIDERS must select at least one provider")
	}
	return providers, nil
}

func buildProviderDefinitions() map[string]providerDefinition {
	definitions := make(map[string]providerDefinition, len(modelsDevProviderIDs)+1)
	for sourceID, catalogID := range modelsDevProviderIDs {
		prefix := catalogID
		if separator := strings.LastIndexByte(prefix, '.'); separator >= 0 {
			prefix = prefix[separator+1:]
		}
		definitions[sourceID] = providerDefinition{
			CatalogID:      catalogID,
			ModelsDevID:    sourceID,
			ModelPrefix:    prefix + "/",
			SourceProvider: sourceID,
		}
	}
	definitions["novita"] = providerDefinition{
		CatalogID:      "custom",
		ModelPrefix:    "novita/",
		SourceProvider: "novita",
		Novita:         true,
	}
	return definitions
}

func syncCatalog(ctx context.Context, providers providerSelection) error {
	namespace := os.Getenv("POD_NAMESPACE")
	if namespace == "" {
		return errors.New("POD_NAMESPACE is required")
	}

	client := &http.Client{Timeout: 30 * time.Second}
	var upstream modelsDevCatalog
	for _, provider := range providers {
		if provider.ModelsDevID != "" {
			if err := fetchJSON(ctx, client, modelsDevURL, "", &upstream); err != nil {
				return fmt.Errorf("fetch models.dev: %w", err)
			}
			break
		}
	}
	var novita novitaCatalog
	if _, enabled := providers["novita"]; enabled {
		novitaAPIKey := os.Getenv("NOVITA_API_KEY")
		if novitaAPIKey == "" {
			return errors.New("NOVITA_API_KEY is required when Novita is enabled")
		}
		if err := fetchJSON(ctx, client, novitaModelsURL, "Bearer "+novitaAPIKey, &novita); err != nil {
			return fmt.Errorf("fetch Novita: %w", err)
		}
	}

	output, err := buildCatalog(upstream, novita, providers)
	if err != nil {
		return err
	}
	encoded, err := json.Marshal(output)
	if err != nil {
		return fmt.Errorf("encode catalog: %w", err)
	}
	return updateConfigMap(ctx, namespace, encoded)
}

func run(ctx context.Context, providers providerSelection) error {
	syncInterval, err := durationFromEnvironment("SYNC_INTERVAL", defaultSyncInterval)
	if err != nil {
		return err
	}
	syncRetryInterval, err := durationFromEnvironment("SYNC_RETRY_INTERVAL", defaultSyncRetryInterval)
	if err != nil {
		return err
	}
	go syncOnSchedule(ctx, syncInterval, syncRetryInterval, providers)
	return serveCatalog(providers)
}

func durationFromEnvironment(name, defaultValue string) (time.Duration, error) {
	value := os.Getenv(name)
	if value == "" {
		value = defaultValue
	}
	duration, err := time.ParseDuration(value)
	if err != nil {
		return 0, fmt.Errorf("parse %s: %w", name, err)
	}
	if duration <= 0 {
		return 0, fmt.Errorf("%s must be positive", name)
	}
	return duration, nil
}

func syncOnSchedule(
	ctx context.Context,
	syncInterval time.Duration,
	syncRetryInterval time.Duration,
	providers providerSelection,
) {
	delay := time.Duration(0)
	for {
		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-timer.C:
		}
		if err := syncCatalog(ctx, providers); err != nil {
			fmt.Fprintln(os.Stderr, err)
			delay = syncRetryInterval
		} else {
			delay = syncInterval
		}
	}
}

func serveCatalog(providers providerSelection) error {
	server := &http.Server{
		Addr:              ":8080",
		Handler:           catalogHandler("/catalog/catalog.json", providers),
		ReadHeaderTimeout: 5 * time.Second,
		IdleTimeout:       60 * time.Second,
	}
	return server.ListenAndServe()
}

func catalogHandler(path string, providers providerSelection) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", func(response http.ResponseWriter, _ *http.Request) {
		response.WriteHeader(http.StatusNoContent)
	})
	mux.HandleFunc("/readyz", func(response http.ResponseWriter, _ *http.Request) {
		input, err := readCatalog(path)
		if err != nil || !catalogReady(input, providers) {
			http.Error(response, "catalog invalid", http.StatusServiceUnavailable)
			return
		}
		response.WriteHeader(http.StatusNoContent)
	})
	mux.HandleFunc("/metrics", func(response http.ResponseWriter, _ *http.Request) {
		input, err := readCatalog(path)
		if err != nil {
			http.Error(response, "catalog invalid", http.StatusServiceUnavailable)
			return
		}
		response.Header().Set("Content-Type", "text/plain; version=0.0.4")
		_, _ = response.Write(renderMetrics(input))
	})
	mux.HandleFunc("/v1/models", func(response http.ResponseWriter, request *http.Request) {
		if request.Method != http.MethodGet {
			response.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		input, err := readCatalog(path)
		if err != nil {
			http.Error(response, "catalog invalid", http.StatusServiceUnavailable)
			return
		}
		response.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(response).Encode(renderModelList(input))
	})
	return mux
}

func catalogReady(input catalog, providers providerSelection) bool {
	for _, provider := range providers {
		if len(input.Providers[provider.CatalogID].Models) == 0 {
			return false
		}
	}
	return true
}

func readCatalog(path string) (catalog, error) {
	contents, err := os.ReadFile(path)
	if err != nil {
		return catalog{}, err
	}
	input := catalog{}
	if err := json.Unmarshal(contents, &input); err != nil {
		return catalog{}, err
	}
	return input, nil
}

func renderMetrics(input catalog) []byte {
	var output strings.Builder
	output.WriteString("# HELP agentgateway_model_token_price_usd_per_million Model token price in USD per million tokens.\n")
	output.WriteString("# TYPE agentgateway_model_token_price_usd_per_million gauge\n")
	providers := sortedKeys(input.Providers)
	for _, providerID := range providers {
		provider := input.Providers[providerID]
		for _, modelName := range sortedKeys(provider.Models) {
			model := provider.Models[modelName]
			for _, tokenType := range sortedKeys(model.Rates) {
				fmt.Fprintf(
					&output,
					"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"%s\",gen_ai_request_model=\"%s\",gen_ai_token_type=\"%s\",source_provider=\"%s\"} %s\n",
					escapeLabel(providerID),
					escapeLabel(modelName),
					escapeLabel(metricTokenType(tokenType)),
					escapeLabel(sourceProvider(providerID)),
					model.Rates[tokenType],
				)
			}
		}
	}
	return []byte(output.String())
}

func renderModelList(input catalog) modelList {
	output := modelList{Object: "list"}
	for _, providerID := range sortedKeys(input.Providers) {
		provider := input.Providers[providerID]
		for _, modelName := range sortedKeys(provider.Models) {
			output.Data = append(output.Data, modelListItem{
				ID:      modelPrefix(providerID) + modelName,
				Object:  "model",
				Created: 0,
				OwnedBy: sourceProvider(providerID),
			})
		}
	}
	return output
}

func sortedKeys[Value any](values map[string]Value) []string {
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}

func escapeLabel(value string) string {
	return strings.NewReplacer("\\", "\\\\", "\n", "\\n", "\"", "\\\"").Replace(value)
}

func metricTokenType(value string) string {
	replacements := map[string]string{
		"cacheRead":   "cache_read",
		"cacheWrite":  "cache_write",
		"inputAudio":  "input_audio",
		"outputAudio": "output_audio",
	}
	if replacement, found := replacements[value]; found {
		return replacement
	}
	return value
}

func sourceProvider(providerID string) string {
	for _, provider := range buildProviderDefinitions() {
		if provider.CatalogID == providerID {
			return provider.SourceProvider
		}
	}
	return providerID
}

func modelPrefix(providerID string) string {
	for _, provider := range buildProviderDefinitions() {
		if provider.CatalogID == providerID {
			return provider.ModelPrefix
		}
	}
	return ""
}

func fetchJSON(ctx context.Context, client *http.Client, url, authorization string, target any) error {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	request.Header.Set("Accept", "application/json")
	if authorization != "" {
		request.Header.Set("Authorization", authorization)
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer func() {
		_ = response.Body.Close()
	}()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, response.Body)
		return fmt.Errorf("unexpected HTTP status %s", response.Status)
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, 16<<20))
	decoder.UseNumber()
	return decoder.Decode(target)
}

func buildCatalog(
	upstream modelsDevCatalog,
	novita novitaCatalog,
	enabledProviders providerSelection,
) (catalog, error) {
	providers := make(map[string]catalogProvider, len(enabledProviders))
	for name, definition := range enabledProviders {
		var models map[string]catalogModel
		if definition.Novita {
			var err error
			models, err = convertNovita(novita)
			if err != nil {
				return catalog{}, fmt.Errorf("convert %s catalog: %w", name, err)
			}
		} else {
			models = convertModelsDev(upstream[definition.ModelsDevID])
		}
		if len(models) == 0 {
			return catalog{}, fmt.Errorf("provider %s returned no priced models", name)
		}
		providers[definition.CatalogID] = catalogProvider{Models: models}
	}
	return catalog{Providers: providers}, nil
}

func convertModelsDev(provider modelsDevProvider) map[string]catalogModel {
	models := make(map[string]catalogModel)
	for name, model := range provider.Models {
		rates := ratesFromModelsDevCost(model.Cost)
		if len(rates) == 0 {
			continue
		}
		converted := catalogModel{Rates: rates}
		for _, tier := range model.Cost.Tiers {
			if tier.Tier.Type != "context" || tier.Tier.Size <= 0 {
				continue
			}
			tierRates := ratesFromModelsDevTier(tier)
			if len(tierRates) > 0 {
				converted.Tiers = append(converted.Tiers, catalogTier{ContextOver: tier.Tier.Size, Rates: tierRates})
			}
		}
		sort.Slice(converted.Tiers, func(left, right int) bool {
			return converted.Tiers[left].ContextOver < converted.Tiers[right].ContextOver
		})
		models[name] = converted
	}
	return models
}

func ratesFromModelsDevCost(cost modelsDevCost) map[string]string {
	return ratesFromNumbers(cost.Input, cost.Output, cost.CacheRead, cost.CacheWrite, cost.Reasoning, cost.InputAudio, cost.OutputAudio)
}

func ratesFromModelsDevTier(tier modelsDevTier) map[string]string {
	return ratesFromNumbers(tier.Input, tier.Output, tier.CacheRead, tier.CacheWrite, tier.Reasoning, tier.InputAudio, tier.OutputAudio)
}

func ratesFromNumbers(input, output, cacheRead, cacheWrite, reasoning, inputAudio, outputAudio *json.Number) map[string]string {
	values := []struct {
		name  string
		value *json.Number
	}{
		{"input", input},
		{"output", output},
		{"cacheRead", cacheRead},
		{"cacheWrite", cacheWrite},
		{"reasoning", reasoning},
		{"inputAudio", inputAudio},
		{"outputAudio", outputAudio},
	}
	rates := make(map[string]string)
	for _, value := range values {
		if value.value != nil {
			rates[value.name] = value.value.String()
		}
	}
	return rates
}

func convertNovita(upstream novitaCatalog) (map[string]catalogModel, error) {
	models := make(map[string]catalogModel)
	for _, model := range upstream.Data {
		if model.ID == "" || model.InputPrice == "" || model.OutputPrice == "" {
			continue
		}
		input, err := tenThousandths(model.InputPrice)
		if err != nil {
			return nil, fmt.Errorf("model %q input price: %w", model.ID, err)
		}
		output, err := tenThousandths(model.OutputPrice)
		if err != nil {
			return nil, fmt.Errorf("model %q output price: %w", model.ID, err)
		}
		models[model.ID] = catalogModel{Rates: map[string]string{"input": input, "output": output}}
	}
	return models, nil
}

func tenThousandths(value json.Number) (string, error) {
	raw := value.String()
	if strings.ContainsAny(raw, ".eE-") {
		return "", fmt.Errorf("expected a non-negative integer, got %q", raw)
	}
	raw = strings.TrimLeft(raw, "0")
	if raw == "" {
		return "0", nil
	}
	if len(raw) <= 4 {
		raw = strings.Repeat("0", 4-len(raw)+1) + raw
	}
	integer := raw[:len(raw)-4]
	fraction := strings.TrimRight(raw[len(raw)-4:], "0")
	if fraction == "" {
		return integer, nil
	}
	return integer + "." + fraction, nil
}

func updateConfigMap(ctx context.Context, namespace string, catalogJSON []byte) error {
	token, err := os.ReadFile(filepath.Join(serviceAccountDir, "token"))
	if err != nil {
		return fmt.Errorf("read service account token: %w", err)
	}
	certificate, err := os.ReadFile(filepath.Join(serviceAccountDir, "ca.crt"))
	if err != nil {
		return fmt.Errorf("read Kubernetes CA: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(certificate) {
		return errors.New("parse Kubernetes CA")
	}
	client := &http.Client{
		Timeout: 30 * time.Second,
		Transport: &http.Transport{TLSClientConfig: &tls.Config{
			MinVersion: tls.VersionTLS12,
			RootCAs:    pool,
		}},
	}
	url := "https://kubernetes.default.svc/api/v1/namespaces/" + namespace + "/configmaps"
	return writeConfigMap(ctx, client, url, token, catalogJSON)
}

func writeConfigMap(ctx context.Context, client *http.Client, collectionURL string, token, catalogJSON []byte) error {
	patch := configMapPatch{
		Data: map[string]string{"catalog.json": string(catalogJSON)},
	}
	itemURL := collectionURL + "/" + configMapName
	status, err := kubernetesRequest(ctx, client, http.MethodPatch, itemURL, token, patch, nil)
	if err != nil {
		return fmt.Errorf("patch ConfigMap: %w", err)
	}
	if status == http.StatusOK {
		return nil
	}
	if status != http.StatusNotFound {
		return fmt.Errorf("patch ConfigMap: unexpected HTTP status %d", status)
	}

	resource := configMapResource{
		APIVersion: "v1",
		Kind:       "ConfigMap",
		Metadata: configMapMetadata{
			Name: configMapName,
			Annotations: map[string]string{
				"kustomize.toolkit.fluxcd.io/prune": "disabled",
			},
			Labels: map[string]string{
				"app.kubernetes.io/instance": "agentgateway",
				"app.kubernetes.io/name":     "model-catalog",
				"app.kubernetes.io/part-of":  "agentgateway",
			},
		},
		Data: patch.Data,
	}
	status, err = kubernetesRequest(ctx, client, http.MethodPost, collectionURL, token, resource, nil)
	if err != nil {
		return fmt.Errorf("create ConfigMap: %w", err)
	}
	if status == http.StatusCreated {
		return nil
	}
	if status != http.StatusConflict {
		return fmt.Errorf("create ConfigMap: unexpected HTTP status %d", status)
	}

	status, err = kubernetesRequest(ctx, client, http.MethodPatch, itemURL, token, patch, nil)
	if err != nil {
		return fmt.Errorf("patch ConfigMap after create conflict: %w", err)
	}
	if status != http.StatusOK {
		return fmt.Errorf("patch ConfigMap after create conflict: unexpected HTTP status %d", status)
	}
	return nil
}

func kubernetesRequest(ctx context.Context, client *http.Client, method, url string, token []byte, body any, target any) (int, error) {
	var encoded io.Reader
	if body != nil {
		contents, err := json.Marshal(body)
		if err != nil {
			return 0, err
		}
		encoded = bytes.NewReader(contents)
	}
	request, err := http.NewRequestWithContext(ctx, method, url, encoded)
	if err != nil {
		return 0, err
	}
	request.Header.Set("Authorization", "Bearer "+string(token))
	request.Header.Set("Accept", "application/json")
	if body != nil {
		contentType := "application/json"
		if method == http.MethodPatch {
			contentType = "application/merge-patch+json"
		}
		request.Header.Set("Content-Type", contentType)
	}
	response, err := client.Do(request)
	if err != nil {
		return 0, err
	}
	defer func() {
		_ = response.Body.Close()
	}()
	if target != nil && response.StatusCode >= 200 && response.StatusCode < 300 {
		if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(target); err != nil {
			return response.StatusCode, err
		}
	} else {
		_, _ = io.Copy(io.Discard, response.Body)
	}
	return response.StatusCode, nil
}