package tool

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

	"tcm-agent/internal/config"
)

// ========================================================================
//  MaxKB 知识库客户端
//  ========================================================================
//  MaxKB 是一个开源的企业级 RAG 知识库平台：
//   - 文档管理（PDF/Word/Markdown 等）
//   - 智能分段 + 向量化（pgvector）
//   - 语义检索 + LLM 增强生成
//   - 兼容 OpenAI 协议
//
//  官方文档：https://maxkb.cn/docs/
//  Docker 一键部署：docker run -d --name maxkb -p 8080:8080 1panel/maxkb
//  ========================================================================

// MaxKBClient MaxKB 知识库客户端
type MaxKBClient struct {
	BaseURL string        `json:"base_url"`
	APIKey  string        `json:"api_key"`
	AppID   string        `json:"app_id"`
	Client  *http.Client `json:"-"`
}

// NewMaxKBClient 创建 MaxKB 客户端
//
// 参数：
//   cfg - MaxKB 配置（BaseURL/APIKey/AppID）
//
// 返回：
//   初始化完成的客户端
func NewMaxKBClient(cfg config.MaxKBConfig) *MaxKBClient {
	return &MaxKBClient{
		BaseURL: cfg.BaseURL,
		APIKey:  cfg.APIKey,
		AppID:   cfg.AppID,
		Client:  &http.Client{Timeout: 60 * time.Second},
	}
}

// Chat 调用 MaxKB 的对话接口（兼容 OpenAI 格式）
//
// 这是最核心的方法：把问题发给 MaxKB，
// 它会自动做 RAG 检索 + LLM 生成，返回融合知识库的回答。
//
// 参数：
//   ctx   - 上下文（支持超时取消）
//   query - 用户问题或检索关键词
//
// 返回：
//   MaxKB 生成的回答（已融合知识库检索结果）
func (c *MaxKBClient) Chat(ctx context.Context, query string) (string, error) {
	body := map[string]any{
		"model": "maxkb-model",
		"messages": []map[string]string{
			{"role": "user", "content": query},
		},
		"stream": false,
	}

	buf, _ := json.Marshal(body)
	url := fmt.Sprintf("%s/api/application/%s/chat/completions", c.BaseURL, c.AppID)

	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(buf))
	if err != nil {
		return "", fmt.Errorf("[MaxKB] 创建请求失败: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.Client.Do(req)
	if err != nil {
		return "", fmt.Errorf("[MaxKB] 调用失败: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		data, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("[MaxKB] 返回错误 %d: %s", resp.StatusCode, string(data))
	}

	var result struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	json.NewDecoder(resp.Body).Decode(&result)

	if len(result.Choices) == 0 {
		return "", fmt.Errorf("[MaxKB] 返回空结果")
	}

	log.Printf("[MaxKB] ✅ 检索成功 | 查询: %.50s...", query)
	return result.Choices[0].Message.Content, nil
}

// Search 仅做知识检索（不生成，返回原始片段）
//
// 适合需要"引用来源"的场景。
// 返回 TopK 条相关文档片段。
func (c *MaxKBClient) Search(ctx context.Context, query string, topK int) ([]string, error) {
	if topK <= 0 {
		topK = 5
	}

	body := map[string]any{
		"query": query,
		"top_k": topK,
	}
	buf, _ := json.Marshal(body)
	url := fmt.Sprintf("%s/api/application/%s/search", c.BaseURL, c.AppID)

	req, _ := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(buf))
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.Client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("[MaxKB] 检索失败: %w", err)
	}
	defer resp.Body.Close()

	var result struct {
		Documents []struct {
			Content string  `json:"content"`
			Score   float64 `json:"score"`
		} `json:"documents"`
	}
	json.NewDecoder(resp.Body).Decode(&result)

	docs := make([]string, 0, len(result.Documents))
	for _, d := range result.Documents {
		docs = append(docs, d.Content)
	}
	return docs, nil
}