package main

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/openai"
	"github.com/tmc/langchaingo/memory"
)

const (
	model     = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
	base_url  = "https://api.siliconflow.cn/v1"
	token     = "sk-udqgogvqgfriqvsehkpxgaejfclvoikbntjeaeijhqnhnviw"
	cotPrompt = `请按照以下步骤分析这道算法题：
	1. 算法判断：判断题目类型和所需算法
	2. 基本思路：说明解题思路和关键点
	3. 框架代码：给出代码框架
	4. 易错解析：指出常见错误和注意事项

	题目：%s

	请按上述格式进行分析：`
)

func main() {
	llm, err := openai.New(
		openai.WithModel(model),
		openai.WithBaseURL(base_url),
		openai.WithToken(token),
	)
	if err != nil {
		log.Fatal(err)
	}

	// [Memory] Init
	chatMemory := memory.NewConversationBuffer()
	ctx := context.Background()
	reader := bufio.NewReader(os.Stdin)

	fmt.Println("Chat Application Started. Type 'exit' to quit.")
	fmt.Println("-----------------------------------------------")

	for {
		fmt.Print("You: ")
		userInput, _ := reader.ReadString('\n')
		userInput = strings.TrimSpace(userInput)

		if userInput == "exit" {
			break
		}

		var fullPrompt string
		if isAlgorithmQuestion(userInput) {
			fullPrompt = fmt.Sprintf(cotPrompt, userInput)
		} else {
			messages, _ := chatMemory.ChatHistory.Messages(ctx)
			// Conversation
			var conversation string
			for _, msg := range messages {
				conversation += msg.GetContent() + "\n"
			}
			fullPrompt = conversation + "Human: " + userInput + "\nAssistant:"
		}

		response, err := llms.GenerateFromSinglePrompt(
			ctx,
			llm,
			fullPrompt,
		)

		if err != nil {
			log.Fatal(err)
		}

		chatMemory.ChatHistory.AddUserMessage(ctx, userInput)
		chatMemory.ChatHistory.AddAIMessage(ctx, response)

		fmt.Println("[AI Response] ", response)
	}

}

func isAlgorithmQuestion(input string) bool {
	keywords := []string{"算法", "编程", "题目", "解题", "代码实现"}
	for _, keyword := range keywords {
		if strings.Contains(input, keyword) {
			return true
		}
	}
	return false
}