// Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
//
// 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.
//
// Requirement: Any integration or derivative work must explicitly attribute
// Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
// documentation or user interface, as detailed in the NOTICE file.

package database

import (
	"os"
	"time"

	"gorm.io/gorm"
)

type ModelParams struct {
	BaseUrl string `json:"base_url"`
	Token   string `json:"token"`
	Model   string `json:"model"`
	Limit   int    `json:"limit"`

	ExtraHeaders map[string]string `json:"extra_headers,omitempty"`
	ExtraBody    map[string]any    `json:"extra_body,omitempty"`
}

// Model 模型表
type Model struct {
	ModelID            string   `gorm:"primaryKey;column:model_id" json:"model_id" yaml:"model_id"`                              // 模型ID
	Username           string   `gorm:"column:username;not null" json:"username" yaml:"-"`                                        // 创建者用户名
	ModelName          string   `gorm:"column:model_name;not null" json:"model_name" yaml:"model_name"`                           // 模型名称
	Token              string   `gorm:"column:token;not null" json:"token" yaml:"token"`                                          // API Token
	BaseURL            string   `gorm:"column:base_url;not null" json:"base_url" yaml:"base_url"`                                 // 基础URL
	Note               string   `gorm:"column:note" json:"note" yaml:"note,omitempty"`                                            // 备注信息
	Limit              int      `gorm:"column:limit" json:"limit" yaml:"limit,omitempty"`
	Default            []string `gorm:"-" json:"default,omitempty" yaml:"default,omitempty"`                                      // 默认字段
	CreatedAt          int64    `gorm:"column:created_at;not null" json:"created_at" yaml:"-"`                                    // 时间戳毫秒级
	UpdatedAt          int64    `gorm:"column:updated_at;not null" json:"updated_at" yaml:"-"`                                    // 时间戳毫秒级

	// ExtraHeaders 随每个模型请求发送；ExtraBody 合并进请求体，键支持点号路径
	ExtraHeaders map[string]string `gorm:"column:extra_headers;serializer:json" json:"extra_headers,omitempty" yaml:"extra_headers,omitempty"`
	ExtraBody    map[string]any    `gorm:"column:extra_body;serializer:json" json:"extra_body,omitempty" yaml:"extra_body,omitempty"`

	// 关联关系
	User User `gorm:"foreignKey:Username" json:"user" yaml:"-"`
}

// ModelStore 模型数据存储
type ModelStore struct {
	db *gorm.DB
}

// NewModelStore 创建新的ModelStore实例
func NewModelStore(db *gorm.DB) *ModelStore {
	return &ModelStore{db: db}
}

// Init 自动迁移模型相关表结构
func (s *ModelStore) Init() error {
	if err := s.db.AutoMigrate(&Model{}); err != nil {
		return err
	}
	// 创建索引优化查询
	return s.db.Exec("CREATE INDEX IF NOT EXISTS idx_models_username_created ON models(username, created_at DESC)").Error
}

// CreateModel 创建模型
func (s *ModelStore) CreateModel(model *Model) error {
	now := time.Now().UnixMilli()
	model.CreatedAt = now
	model.UpdatedAt = now
	return s.db.Create(model).Error
}

// GetModel 获取模型信息
func (s *ModelStore) GetModel(modelID string) (*Model, error) {
	var model Model
	err := s.db.Preload("User").First(&model, "model_id = ?", modelID).Error
	if err != nil {
		// Try YAML model
		if yamlModel := s.GetYamlModel(modelID); yamlModel != nil {
			return yamlModel, nil
		}
		return nil, err
	}
	return &model, nil
}

// GetModelByUser 获取用户创建的模型
func (s *ModelStore) GetModelByUser(modelID string, username string) (*Model, error) {
	var model Model
	err := s.db.Preload("User").First(&model, "model_id = ? AND username = ?", modelID, username).Error
	if err != nil {
		return nil, err
	}
	return &model, nil
}

// GetAllModels 获取所有模型
func (s *ModelStore) GetAllModels() ([]*Model, error) {
	var models []*Model
	err := s.db.Preload("User").Order("created_at DESC").Find(&models).Error
	if err != nil {
		return nil, err
	}
	return models, nil
}

// GetUserModels 获取用户的所有模型
func (s *ModelStore) GetUserModels(username string) ([]*Model, error) {
	var models []*Model
	err := s.db.Preload("User").Where("username = ? or username = '' or username = 'public_user'", username).Order("created_at DESC").Find(&models).Error
	if err != nil {
		return nil, err
	}

	// Append YAML models
	yamlModels, _ := s.LoadYamlModels()
	if len(yamlModels) > 0 {
		models = append(models, yamlModels...)
	}

	return models, nil
}

// UpdateModel 更新模型信息
func (s *ModelStore) UpdateModel(modelID string, username string, updates map[string]interface{}) error {
	now := time.Now().UnixMilli()

	// extra_headers / extra_body 是 map 字段：map 形式的 Updates 不会走字段
	// 序列化（serializer:json），驱动会直接拒绝 map 值并报
	// "unsupported type map[string]interface {}"。这两个字段改用
	// struct + Select 写入，空 map 也能落库（即清空）。
	extra := &Model{UpdatedAt: now}
	selected := []string{"updated_at"}
	if value, ok := updates["extra_headers"]; ok {
		extra.ExtraHeaders, _ = value.(map[string]string)
		selected = append(selected, "extra_headers")
		delete(updates, "extra_headers")
	}
	if value, ok := updates["extra_body"]; ok {
		extra.ExtraBody, _ = value.(map[string]any)
		selected = append(selected, "extra_body")
		delete(updates, "extra_body")
	}

	if len(selected) > 1 {
		err := s.db.Model(&Model{}).
			Where("model_id = ? AND username = ?", modelID, username).
			Select(selected).
			Updates(extra).Error
		if err != nil {
			return err
		}
	}

	// 添加更新时间
	updates["updated_at"] = now
	return s.db.Model(&Model{}).Where("model_id = ? AND username = ?", modelID, username).Updates(updates).Error
}

// DeleteModel 删除模型
func (s *ModelStore) DeleteModel(modelID string, username string) error {
	return s.db.Delete(&Model{}, "model_id = ? AND username = ?", modelID, username).Error
}

// BatchDeleteModels 批量删除模型
func (s *ModelStore) BatchDeleteModels(modelIDs []string, username string) (int64, error) {
	result := s.db.Delete(&Model{}, "model_id IN ? AND username = ?", modelIDs, username)
	return result.RowsAffected, result.Error
}

// CheckModelExists 检查模型是否存在
func (s *ModelStore) CheckModelExists(modelID string) (bool, error) {
	var count int64
	err := s.db.Model(&Model{}).Where("model_id = ?", modelID).Count(&count).Error
	if count > 0 {
		return true, nil
	}
	// Check YAML
	if s.GetYamlModel(modelID) != nil {
		return true, nil
	}
	return false, err
}

// CheckModelExistsByUser 检查用户是否拥有该模型
func (s *ModelStore) CheckModelExistsByUser(modelID string, username string) (bool, error) {
	var count int64
	err := s.db.Model(&Model{}).Where("model_id = ? AND username = ?", modelID, username).Count(&count).Error
	return count > 0, err
}

func (s *ModelStore) AutoAddModels() {
	// 判断如果模型为空，并且环境变量存在 model token base_url，则自动添加模型
	if s.db == nil {
		return
	}
	var count int64
	s.db.Model(&Model{}).Count(&count)
	if count == 0 {
		model := os.Getenv("MODEL")
		token := os.Getenv("TOKEN")
		baseUrl := os.Getenv("BASE_URL")
		if model != "" && token != "" && baseUrl != "" {
			s.CreateModel(&Model{
				ModelID:   "system_default",
				Username:  "",
				ModelName: model,
				Token:     token,
				BaseURL:   baseUrl,
				Note:      "系统默认内置",
				CreatedAt: time.Now().UnixMilli(),
				UpdatedAt: time.Now().UnixMilli(),
			})
		}
	}
}