package controller

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/songquanpeng/one-api/common/ctxkey"
	"github.com/songquanpeng/one-api/model"
	relay "github.com/songquanpeng/one-api/relay"
	"github.com/songquanpeng/one-api/relay/adaptor/openai"
	"github.com/songquanpeng/one-api/relay/apitype"
	"github.com/songquanpeng/one-api/relay/channeltype"
	"github.com/songquanpeng/one-api/relay/meta"
	relaymodel "github.com/songquanpeng/one-api/relay/model"
)

// https://platform.openai.com/docs/api-reference/models/list

type OpenAIModelPermission struct {
	Id                 string  `json:"id"`
	Object             string  `json:"object"`
	Created            int     `json:"created"`
	AllowCreateEngine  bool    `json:"allow_create_engine"`
	AllowSampling      bool    `json:"allow_sampling"`
	AllowLogprobs      bool    `json:"allow_logprobs"`
	AllowSearchIndices bool    `json:"allow_search_indices"`
	AllowView          bool    `json:"allow_view"`
	AllowFineTuning    bool    `json:"allow_fine_tuning"`
	Organization       string  `json:"organization"`
	Group              *string `json:"group"`
	IsBlocking         bool    `json:"is_blocking"`
}

type OpenAIModels struct {
	Id         string                  `json:"id"`
	Object     string                  `json:"object"`
	Created    int                     `json:"created"`
	OwnedBy    string                  `json:"owned_by"`
	Permission []OpenAIModelPermission `json:"permission"`
	Root       string                  `json:"root"`
	Parent     *string                 `json:"parent"`
}

var models []OpenAIModels
var modelsMap map[string]OpenAIModels
var channelId2Models map[int][]string

func init() {
	var permission []OpenAIModelPermission
	permission = append(permission, OpenAIModelPermission{
		Id:                 "modelperm-LwHkVFn8AcMItP432fKKDIKJ",
		Object:             "model_permission",
		Created:            1626777600,
		AllowCreateEngine:  true,
		AllowSampling:      true,
		AllowLogprobs:      true,
		AllowSearchIndices: false,
		AllowView:          true,
		AllowFineTuning:    false,
		Organization:       "*",
		Group:              nil,
		IsBlocking:         false,
	})
	// https://platform.openai.com/docs/models/model-endpoint-compatibility
	for i := 0; i < apitype.Dummy; i++ {
		if i == apitype.AIProxyLibrary {
			continue
		}
		adaptor := relay.GetAdaptor(i)
		channelName := adaptor.GetChannelName()
		modelNames := adaptor.GetModelList()
		for _, modelName := range modelNames {
			models = append(models, OpenAIModels{
				Id:         modelName,
				Object:     "model",
				Created:    1626777600,
				OwnedBy:    channelName,
				Permission: permission,
				Root:       modelName,
				Parent:     nil,
			})
		}
	}
	for _, channelType := range openai.CompatibleChannels {
		if channelType == channeltype.Azure {
			continue
		}
		channelName, channelModelList := openai.GetCompatibleChannelMeta(channelType)
		for _, modelName := range channelModelList {
			models = append(models, OpenAIModels{
				Id:         modelName,
				Object:     "model",
				Created:    1626777600,
				OwnedBy:    channelName,
				Permission: permission,
				Root:       modelName,
				Parent:     nil,
			})
		}
	}
	modelsMap = make(map[string]OpenAIModels)
	for _, model := range models {
		modelsMap[model.Id] = model
	}
	channelId2Models = make(map[int][]string)
	for i := 1; i < channeltype.Dummy; i++ {
		adaptor := relay.GetAdaptor(channeltype.ToAPIType(i))
		meta := &meta.Meta{
			ChannelType: i,
		}
		adaptor.Init(meta)
		channelId2Models[i] = adaptor.GetModelList()
	}
}

func DashboardListModels(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"message": "",
		"data":    channelId2Models,
	})
}

func ListAllModels(c *gin.Context) {
	c.JSON(200, gin.H{
		"object": "list",
		"data":   models,
	})
}

func ListModels(c *gin.Context) {
	ctx := c.Request.Context()
	var availableModels []string
	if c.GetString(ctxkey.AvailableModels) != "" {
		availableModels = strings.Split(c.GetString(ctxkey.AvailableModels), ",")
	} else {
		userId := c.GetInt(ctxkey.Id)
		userGroup, _ := model.CacheGetUserGroup(userId)
		availableModels, _ = model.CacheGetGroupModels(ctx, userGroup)
	}
	modelSet := make(map[string]bool)
	for _, availableModel := range availableModels {
		modelSet[availableModel] = true
	}
	availableOpenAIModels := make([]OpenAIModels, 0)
	for _, model := range models {
		if _, ok := modelSet[model.Id]; ok {
			modelSet[model.Id] = false
			availableOpenAIModels = append(availableOpenAIModels, model)
		}
	}
	for modelName, ok := range modelSet {
		if ok {
			availableOpenAIModels = append(availableOpenAIModels, OpenAIModels{
				Id:      modelName,
				Object:  "model",
				Created: 1626777600,
				OwnedBy: "custom",
				Root:    modelName,
				Parent:  nil,
			})
		}
	}
	c.JSON(200, gin.H{
		"object": "list",
		"data":   availableOpenAIModels,
	})
}

func RetrieveModel(c *gin.Context) {
	modelId := c.Param("model")
	if model, ok := modelsMap[modelId]; ok {
		c.JSON(200, model)
	} else {
		Error := relaymodel.Error{
			Message: fmt.Sprintf("The model '%s' does not exist", modelId),
			Type:    "invalid_request_error",
			Param:   "model",
			Code:    "model_not_found",
		}
		c.JSON(200, gin.H{
			"error": Error,
		})
	}
}

func GetUserAvailableModels(c *gin.Context) {
	ctx := c.Request.Context()
	id := c.GetInt(ctxkey.Id)
	userGroup, err := model.CacheGetUserGroup(id)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": err.Error(),
		})
		return
	}
	models, err := model.CacheGetGroupModels(ctx, userGroup)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": err.Error(),
		})
		return
	}
	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"message": "",
		"data":    models,
	})
	return
}

// --- Model Catalog (powered by Zenmux) ---

// ModelCatalogItem is the enriched model record returned to aiscri-xiong clients.
type ModelCatalogItem struct {
	Id               string   `json:"id"`
	Name             string   `json:"name"`
	Provider         string   `json:"provider"`
	InputModalities  []string `json:"input_modalities"`
	OutputModalities []string `json:"output_modalities"`
	ContextLength    int      `json:"context_length"`
	MaxOutputTokens  int      `json:"max_output_tokens"`
	InputPrice       float64  `json:"input_price"`
	OutputPrice      float64  `json:"output_price"`
	Tags             []string `json:"tags"`
	// SupportsReasoning: 0=none, 1=always-on, 2=toggleable
	SupportsReasoning int    `json:"supports_reasoning"`
	SuitableApi       string `json:"suitable_api"`
}

// zenmuxModel is the raw response shape from Zenmux listByFilter API.
type zenmuxModel struct {
	Slug                string  `json:"slug"`
	Name                string  `json:"name"`
	Author              string  `json:"author"`
	InputModalities     string  `json:"input_modalities"`  // comma-separated
	OutputModalities    string  `json:"output_modalities"` // comma-separated
	ContextLength       int     `json:"context_length"`
	MaxCompletionTokens int     `json:"max_completion_tokens"`
	PricingPrompt       string  `json:"pricing_prompt"`
	PricingCompletion   string  `json:"pricing_completion"`
	SuitableApi         string  `json:"suitable_api"` // comma-separated
	SupportsReasoning   int     `json:"supports_reasoning"`
	SupportedParameters string  `json:"supported_parameters"` // comma-separated
}

type zenmuxListResponse struct {
	Success bool          `json:"success"`
	Data    []zenmuxModel `json:"data"`
}

var (
	catalogCache     []ModelCatalogItem
	catalogCacheTime time.Time
	catalogMu        sync.RWMutex
	catalogCacheTTL  = 5 * time.Minute

	// zenmuxCToken is the public frontend token embedded in the Zenmux website.
	// Override via env var ZENMUX_CTOKEN if it ever changes.
	zenmuxCToken = func() string {
		if v := os.Getenv("ZENMUX_CTOKEN"); v != "" {
			return v
		}
		return "2uF44yawHUs41edv2fRcV_eE"
	}()
)

func splitCSV(s string) []string {
	if s == "" {
		return nil
	}
	parts := strings.Split(s, ",")
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		if t := strings.TrimSpace(p); t != "" {
			out = append(out, t)
		}
	}
	return out
}

func parsePrice(s string) float64 {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0
	}
	var f float64
	fmt.Sscanf(s, "%f", &f)
	return f
}

func zenmuxToItem(z zenmuxModel) ModelCatalogItem {
	inMod := splitCSV(z.InputModalities)
	outMod := splitCSV(z.OutputModalities)
	params := splitCSV(z.SupportedParameters)

	// Derive tags
	tags := []string{}
	for _, p := range params {
		if p == "tools" {
			tags = append(tags, "function-call")
			break
		}
	}
	for _, m := range inMod {
		if m == "image" {
			tags = append(tags, "vision")
			break
		}
	}
	if z.SupportsReasoning > 0 {
		tags = append(tags, "reasoning")
	}

	return ModelCatalogItem{
		Id:                z.Slug,
		Name:              z.Name,
		Provider:          z.Author,
		InputModalities:   inMod,
		OutputModalities:  outMod,
		ContextLength:     z.ContextLength,
		MaxOutputTokens:   z.MaxCompletionTokens,
		InputPrice:        parsePrice(z.PricingPrompt),
		OutputPrice:       parsePrice(z.PricingCompletion),
		Tags:              tags,
		SupportsReasoning: z.SupportsReasoning,
		SuitableApi:       z.SuitableApi,
	}
}

func fetchZenmuxCatalog() ([]ModelCatalogItem, error) {
	catalogMu.RLock()
	if catalogCache != nil && time.Since(catalogCacheTime) < catalogCacheTTL {
		defer catalogMu.RUnlock()
		return catalogCache, nil
	}
	catalogMu.RUnlock()

	url := "https://zenmux.ai/api/frontend/model/listByFilter?ctoken=" + zenmuxCToken + "&sort=newest&context_length=&keyword="
	client := &http.Client{Timeout: 15 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		return nil, fmt.Errorf("zenmux catalog fetch failed: %w", err)
	}
	defer resp.Body.Close()

	var parsed zenmuxListResponse
	if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
		return nil, fmt.Errorf("zenmux catalog parse failed: %w", err)
	}
	if !parsed.Success {
		return nil, fmt.Errorf("zenmux catalog returned success=false")
	}

	items := make([]ModelCatalogItem, 0, len(parsed.Data))
	for _, z := range parsed.Data {
		if z.Slug != "" {
			items = append(items, zenmuxToItem(z))
		}
	}

	catalogMu.Lock()
	catalogCache = items
	catalogCacheTime = time.Now()
	catalogMu.Unlock()

	return items, nil
}

// GetModelCatalog returns models with capability metadata, filtered to only those
// available to the requesting user's group.
func GetModelCatalog(c *gin.Context) {
	ctx := c.Request.Context()
	id := c.GetInt(ctxkey.Id)

	userGroup, err := model.CacheGetUserGroup(id)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
		return
	}

	availableModels, err := model.CacheGetGroupModels(ctx, userGroup)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
		return
	}

	// Build a set for O(1) lookup
	available := make(map[string]bool, len(availableModels))
	for _, m := range availableModels {
		available[m] = true
	}

	catalog, err := fetchZenmuxCatalog()
	if err != nil {
		// Return empty catalog with warning rather than failing hard
		c.JSON(http.StatusOK, gin.H{
			"success": true,
			"message": "model catalog unavailable: " + err.Error(),
			"data":    []ModelCatalogItem{},
		})
		return
	}

	// Filter to models actually available to this user
	result := make([]ModelCatalogItem, 0, len(catalog))
	for _, item := range catalog {
		if available[item.Id] {
			result = append(result, item)
		}
	}

	c.JSON(http.StatusOK, gin.H{
		"success": true,
		"message": "",
		"data":    result,
	})
}