package knowledge_base

import (
	"encoding/json"
	"fmt"
	"strings"
)

type openAIChunk struct {
	Choices []struct {
		Delta        *Delta   `json:"delta"`
		Message      *Message `json:"message"`
		FinishReason *string  `json:"finish_reason"`
	} `json:"choices"`
}

type Delta struct {
	ReasoningContent string `json:"reasoning_content"` // 推理内容
	Content          string `json:"content"`           // 内容
	XRagStatus       string `json:"x_rag_status"`      // rag命中状态 hit|miss
}

type Message struct {
	Role       string `json:"role"`         // 角色
	Content    string `json:"content"`      // 内容
	XRagStatus string `json:"x_rag_status"` // rag命中状态 hit|miss
}

func ParseOpenAIStreamData(dataLine string) (delta *Delta, done bool, err error) {
	data := strings.TrimSpace(strings.TrimPrefix(dataLine, "data:"))
	if data == "" {
		return nil, false, nil
	}

	data = strings.TrimSpace(data)
	if data == "" {
		return nil, false, nil
	}
	if data == "[DONE]" {
		return nil, true, nil
	}

	var chunk openAIChunk
	if err := json.Unmarshal([]byte(data), &chunk); err != nil {
		return nil, false, fmt.Errorf("unmarshal openai stream chunk failed: %w", err)
	}

	for _, c := range chunk.Choices {
		if c.Delta != nil {
			return c.Delta, false, nil // 只输出第一个delta
		}
	}

	return nil, false, nil
}

func ParseOpenAIHTTPData(body string) (message *Message, done bool, err error) {
	data := strings.TrimSpace(body)
	if data == "" {
		return nil, false, nil
	}

	var resp openAIChunk
	if err := json.Unmarshal([]byte(data), &resp); err != nil {
		return nil, false, fmt.Errorf("unmarshal openai stream chunk failed: %w", err)
	}

	for _, c := range resp.Choices {
		if c.Message != nil {
			return c.Message, true, nil // 只输出第一个message
		}
	}

	return nil, false, nil
}