package completion

import (
	"context"
	"errors"
	"io"
	"strings"
	"time"

	"aichat/llm"
	"aichat/message"
	"aichat/stream"
	"aichat/utils"

	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
)

type ChatCompleter func(context.Context, *llm.Profile, model.CreateChatCompletionRequest, time.Duration) (model.ChatCompletionResponse, error)

func CompleteText(ctx context.Context, profile *llm.Profile, chatMessages []message.ChatMessage, maxTokens int) (string, error) {
	return CompleteTextWithTimeout(ctx, profile, chatMessages, maxTokens, time.Duration(profile.Config.Timeout)*time.Second)
}

func CompleteTextWithTimeout(ctx context.Context, profile *llm.Profile, chatMessages []message.ChatMessage, maxTokens int, timeout time.Duration) (string, error) {
	messages, err := message.BuildArkMessages(chatMessages)
	if err != nil {
		return "", err
	}
	completionCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	streamResp, err := profile.Client.CreateChatCompletionStream(completionCtx, model.CreateChatCompletionRequest{
		Model:     profile.Config.Model,
		Messages:  messages,
		MaxTokens: utils.IntPtr(maxTokens),
	}.WithStream(true))
	if err != nil {
		return "", err
	}
	defer streamResp.Close()

	promptTokens := stream.EstimateChatMessagesTokens(chatMessages)
	completionTokens := 0
	parseThinkTags := llm.ShouldParseThinkTags(profile)
	thinkParser := &stream.Parser{}
	var b strings.Builder
	appendVisible := func(delta string) {
		if delta == "" {
			return
		}
		b.WriteString(delta)
		completionTokens += stream.EstimateTokenCount(delta)
	}
	for {
		resp, err := streamResp.Recv()
		if errors.Is(err, io.EOF) {
			if parseThinkTags {
				visible, _ := thinkParser.Flush()
				appendVisible(visible)
			}
			if tracker := stream.TrackerFromContext(ctx); tracker != nil {
				tracker.AddTool(promptTokens, completionTokens)
			}
			return b.String(), nil
		}
		if err != nil {
			return "", err
		}
		if len(resp.Choices) > 0 {
			delta := resp.Choices[0].Delta.Content
			if parseThinkTags {
				visible, _ := thinkParser.Accept(delta)
				appendVisible(visible)
			} else {
				appendVisible(delta)
			}
		}
	}
}

func CompleteChatWithTimeout(ctx context.Context, profile *llm.Profile, request model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
	completionCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	return profile.Client.CreateChatCompletion(completionCtx, request.WithStream(false))
}