package ai

import (
	"context"
	"net/url"
	"strings"
	"time"

	openaiModel "github.com/cloudwego/eino-ext/components/model/openai"
	"github.com/cloudwego/eino/components/model"
	"github.com/gogf/gf/v2/errors/gerror"
)

type OpenAICompatibleConfig struct {
	APIKey     string
	BaseURL    string
	Endpoint   string
	Model      string
	Timeout    time.Duration
	MaxTokens  int
	JSONOutput bool
}

const ProviderOpenAICompatible = "openai-compatible"

type ChatModelConfig struct {
	Provider   string
	APIKey     string
	BaseURL    string
	Endpoint   string
	Model      string
	Timeout    time.Duration
	MaxTokens  int
	JSONOutput bool
}

func NewChatModel(ctx context.Context, config ChatModelConfig) (model.BaseChatModel, error) {
	provider := strings.ToLower(strings.TrimSpace(config.Provider))
	if provider == "" {
		provider = ProviderOpenAICompatible
	}
	switch provider {
	case ProviderOpenAICompatible, "openai":
		return NewOpenAICompatibleChatModel(ctx, OpenAICompatibleConfig{
			APIKey:     config.APIKey,
			BaseURL:    config.BaseURL,
			Endpoint:   config.Endpoint,
			Model:      config.Model,
			Timeout:    config.Timeout,
			MaxTokens:  config.MaxTokens,
			JSONOutput: config.JSONOutput,
		})
	default:
		return nil, gerror.Newf("暂不支持AI模型提供方：%s", provider)
	}
}

func NewOpenAICompatibleChatModel(ctx context.Context, config OpenAICompatibleConfig) (model.BaseChatModel, error) {
	apiKey := strings.TrimSpace(config.APIKey)
	if apiKey == "" {
		return nil, gerror.New("AI模型API Key不能为空")
	}
	modelName := strings.TrimSpace(config.Model)
	if modelName == "" {
		return nil, gerror.New("AI模型名称不能为空")
	}
	baseURL, err := ResolveOpenAIBaseURL(config.BaseURL, config.Endpoint)
	if err != nil {
		return nil, err
	}

	modelConfig := &openaiModel.ChatModelConfig{
		APIKey:  apiKey,
		BaseURL: baseURL,
		Model:   modelName,
		Timeout: config.Timeout,
	}
	if config.MaxTokens > 0 {
		modelConfig.MaxTokens = &config.MaxTokens
	}
	if config.JSONOutput {
		modelConfig.ResponseFormat = &openaiModel.ChatCompletionResponseFormat{
			Type: openaiModel.ChatCompletionResponseFormatTypeJSONObject,
		}
	}

	chatModel, err := openaiModel.NewChatModel(ctx, modelConfig)
	if err != nil {
		return nil, gerror.Wrap(err, "初始化Eino AI模型失败")
	}
	return chatModel, nil
}

func ResolveOpenAIBaseURL(baseURL string, endpoint string) (string, error) {
	value := strings.TrimSpace(baseURL)
	if value == "" {
		value = strings.TrimSpace(endpoint)
	}
	if value == "" {
		return "", gerror.New("AI模型BaseURL不能为空")
	}

	parsed, err := url.Parse(value)
	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
		return "", gerror.New("AI模型BaseURL格式不正确")
	}
	parsed.Path = strings.TrimSuffix(strings.TrimRight(parsed.Path, "/"), "/chat/completions")
	parsed.RawPath = ""
	parsed.RawQuery = ""
	parsed.Fragment = ""
	return strings.TrimRight(parsed.String(), "/"), nil
}