package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gcfg"
	"github.com/gogf/gf/v2/os/gctx"
)

type Dataset struct {
	Id   string `json:"id"`
	Name string `json:"name"`
}

type ListDatasetsRes struct {
	Code int       `json:"code"`
	Data []Dataset `json:"data"`
}

type DeleteReq struct {
	Ids []string `json:"ids"`
}

type CommonResponse struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

func main() {
	ctx := gctx.New()

	fmt.Println("🚀 RAGFlow知识库清理工具（简化版）")
	fmt.Println("=" + strings.Repeat("=", 50))

	// 设置配置文件路径
	g.Cfg().GetAdapter().(*gcfg.AdapterFile).SetPath("../")

	// 读取配置
	baseURL := g.Cfg().MustGet(ctx, "ragflow.base_url").String()
	apiKey := g.Cfg().MustGet(ctx, "ragflow.api_key").String()

	if baseURL == "" || apiKey == "" {
		fmt.Println("❌ RAGFlow配置缺失，请检查../config.yaml")
		fmt.Println("   需要配置: ragflow.base_url 和 ragflow.api_key")
		os.Exit(1)
	}

	baseURL = strings.TrimSuffix(baseURL, "/")
	fmt.Printf("📡 连接到RAGFlow: %s\n", baseURL)

	// 创建HTTP客户端
	client := &http.Client{
		Timeout: 30 * time.Second,
	}

	// 1. 列出所有知识库
	listURL := baseURL + "/api/v1/datasets?page=1&page_size=1000"
	req, err := http.NewRequest("GET", listURL, nil)
	if err != nil {
		fmt.Printf("❌ 创建请求失败: %v\n", err)
		os.Exit(1)
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)

	resp, err := client.Do(req)
	if err != nil {
		fmt.Printf("❌ 查询知识库列表失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	var listRes ListDatasetsRes
	if err := json.Unmarshal(body, &listRes); err != nil {
		fmt.Printf("❌ 解析响应失败: %v\n", err)
		fmt.Printf("响应内容: %s\n", string(body))
		os.Exit(1)
	}

	if len(listRes.Data) == 0 {
		fmt.Println("✅ 没有找到任何知识库，无需清理")
		os.Exit(0)
	}

	// 显示知识库列表
	fmt.Printf("\n📚 发现 %d 个知识库：\n", len(listRes.Data))
	for i, dataset := range listRes.Data {
		fmt.Printf("  %d. ID: %s, Name: %s\n", i+1, dataset.Id, dataset.Name)
	}

	// 二次确认
	fmt.Printf("\n⚠️  警告：即将删除所有 %d 个知识库！\n", len(listRes.Data))
	fmt.Print("请输入 'YES' 确认删除: ")

	var confirm string
	fmt.Scanln(&confirm)

	if confirm != "YES" {
		fmt.Println("❌ 取消删除操作")
		os.Exit(0)
	}

	// 2. 批量删除
	var datasetIds []string
	for _, dataset := range listRes.Data {
		datasetIds = append(datasetIds, dataset.Id)
	}

	deleteReq := DeleteReq{Ids: datasetIds}
	reqBody, _ := json.Marshal(deleteReq)

	deleteURL := baseURL + "/api/v1/datasets"
	delReq, err := http.NewRequest("DELETE", deleteURL, bytes.NewBuffer(reqBody))
	if err != nil {
		fmt.Printf("❌ 创建删除请求失败: %v\n", err)
		os.Exit(1)
	}
	delReq.Header.Set("Authorization", "Bearer "+apiKey)
	delReq.Header.Set("Content-Type", "application/json")

	fmt.Println("\n🗑️  开始删除知识库...")
	delResp, err := client.Do(delReq)
	if err != nil {
		fmt.Printf("❌ 删除失败: %v\n", err)
		os.Exit(1)
	}
	defer delResp.Body.Close()

	delBody, _ := io.ReadAll(delResp.Body)
	var commonRes CommonResponse
	if err := json.Unmarshal(delBody, &commonRes); err != nil {
		fmt.Printf("❌ 解析删除响应失败: %v\n", err)
		fmt.Printf("响应内容: %s\n", string(delBody))
		os.Exit(1)
	}

	if commonRes.Code != 0 {
		fmt.Printf("❌ 删除失败: %s\n", commonRes.Message)
		os.Exit(1)
	}

	fmt.Printf("✅ 成功删除 %d 个知识库！\n", len(datasetIds))
	fmt.Println("=" + strings.Repeat("=", 50))
}