package utils_ollama

import (
	"ai_scheduler/internal/config"
	"ai_scheduler/internal/entitys"
	"context"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

	"github.com/ollama/ollama/api"
)

// Client Ollama客户端适配器
type Client struct {
	client *api.Client
	config *config.OllamaConfig
}

// NewClient 创建新的Ollama客户端
func NewClient(config *config.Config) (client *Client, cleanFunc func(), err error) {
	client = &Client{
		config: &config.Ollama,
	}
	url, err := client.getUrl()
	if err != nil {
		return
	}
	client.client = api.NewClient(url, http.DefaultClient)

	cleanup := func() {
		if client != nil {
			client = nil
		}
	}

	return client, cleanup, nil
}

// ToolSelect 工具选择
func (c *Client) ToolSelect(ctx context.Context, messages []api.Message, tools []api.Tool) (res api.ChatResponse, err error) {

	// 构建聊天请求
	req := &api.ChatRequest{
		Model:     c.config.Model,
		Messages:  messages,
		Stream:    new(bool), // 设置为false，不使用流式响应
		Think:     &api.ThinkValue{Value: false},
		Tools:     tools,
		KeepAlive: &api.Duration{Duration: 24 * time.Hour},
	}
	err = c.client.Chat(ctx, req, func(resp api.ChatResponse) error {
		res = resp
		return nil
	})

	if err != nil {
		return
	}

	return
}

func (c *Client) ChatStream(ctx context.Context, ch chan entitys.Response, messages []api.Message, index string, model string) (err error) {
	if len(model) == 0 {
		model = c.config.Model
	}
	// 构建聊天请求
	req := &api.ChatRequest{
		Model:    model,
		Messages: messages,
		Stream:   nil,
		Think:    &api.ThinkValue{Value: true},
	}
	var w sync.WaitGroup
	w.Add(1)
	go func() {
		defer w.Done()
		err = c.client.Chat(ctx, req, func(resp api.ChatResponse) error {
			if resp.Message.Content != "" {
				entitys.ResStream(ch, index, resp.Message.Content)
			}
			return nil
		})
		if err != nil {
			return
		}
	}()
	w.Wait()

	return
}

func (c *Client) Chat(ctx context.Context, model string, messages []api.Message) (res api.ChatResponse, err error) {
	// 构建聊天请求
	req := &api.ChatRequest{
		Model:    model,
		Messages: messages,
		Stream:   new(bool), // 设置为false，不使用流式响应
		Think:    &api.ThinkValue{Value: false},
	}
	err = c.client.Chat(ctx, req, func(resp api.ChatResponse) error {
		res = resp
		return nil
	})
	if err != nil {
		return
	}

	return
}

func (c *Client) Generation(ctx context.Context, generateRequest *api.GenerateRequest) (result api.GenerateResponse, err error) {
	err = c.client.Generate(ctx, generateRequest, func(resp api.GenerateResponse) error {
		result = resp
		return nil
	})
	return
}

// convertResponse 转换响应格式
func (c *Client) convertResponse(resp *api.ChatResponse) *entitys.ChatResponse {
	//result := &entitys.ChatResponse{
	//	Message:  resp.Message.Content,
	//	Finished: resp.Done,
	//}
	//
	//// 转换工具调用
	//if len(resp.Message.ToolCalls) > 0 {
	//	result.ToolCalls = make([]entitys.ToolCall, len(resp.Message.ToolCalls))
	//	for i, toolCall := range resp.Message.ToolCalls {
	//		// 转换函数参数
	//		argBytes, _ := json.Marshal(toolCall.Function.Arguments)
	//
	//		result.ToolCalls[i] = entitys.ToolCall{
	//			ID:   fmt.Sprintf("call_%d", i),
	//			Type: "function",
	//			Function: entitys.FunctionCall{
	//				Name:      toolCall.Function.Name,
	//				Arguments: json.RawMessage(argBytes),
	//			},
	//		}
	//	}
	//}

	//return result
	return nil
}

func (c *Client) getUrl() (*url.URL, error) {
	baseURL := c.config.BaseURL
	envURL := os.Getenv("OLLAMA_BASE_URL")
	if envURL != "" {
		baseURL = envURL
	}

	return url.Parse(baseURL)
}
