package models

import "strings"

type ModelRegistry struct {
	Providers map[string]ProviderInfo `json:"-"`
}

type ProviderInfo struct {
	ID     string                   `json:"id"`
	Name   string                   `json:"name"`
	Models map[string]ModelMetadata `json:"models"`
}

type ModelMetadata struct {
	ID         string           `json:"id"`
	Name       string           `json:"name"`
	Cost       *ModelCost       `json:"cost,omitempty"`
	Limit      *ModelLimit      `json:"limit,omitempty"`
	Modalities *ModelModalities `json:"modalities,omitempty"`
	ToolCall   *bool            `json:"tool_call,omitempty"`
	Reasoning  *bool            `json:"reasoning,omitempty"`
}

type ModelCost struct {
	Input      float64  `json:"input"`
	Output     float64  `json:"output"`
	CacheRead  *float64 `json:"cache_read,omitempty"`
	CacheWrite *float64 `json:"cache_write,omitempty"`
}

type ModelLimit struct {
	Context uint32 `json:"context"`
	Output  uint32 `json:"output"`
}

type ModelModalities struct {
	Input  []string `json:"input"`
	Output []string `json:"output"`
}

func (r *ModelRegistry) FindModel(modelID string) *ModelMetadata {
	// First try exact match in models map
	for _, provider := range r.Providers {
		if model, ok := provider.Models[modelID]; ok {
			return &model
		}
	}

	// Try searching by ID in metadata
	for _, provider := range r.Providers {
		for _, model := range provider.Models {
			if model.ID == modelID {
				return &model
			}
		}
	}

	// Try fuzzy matching (e.g. gpt-4o-2024-05-13 matching gpt-4o)
	for _, provider := range r.Providers {
		for id, model := range provider.Models {
			if strings.HasPrefix(modelID, id) {
				return &model
			}
		}
	}

	return nil
}