package quickstart

import (
	"context"
	"fmt"
	"os"

	"github.com/bububa/instructor-go"

	"github.com/bububa/atomic-agents/agents"
	"github.com/bububa/atomic-agents/components"
	"github.com/bububa/atomic-agents/components/systemprompt/cot"
	"github.com/bububa/atomic-agents/examples"
	"github.com/bububa/atomic-agents/schema"
)

// CustomOutput represents the response generated by the chat agent, including suggested follow-up questions.
type CustomOutput struct {
	schema.Base
	// ChatMessage is the chat message exchanged between the user and the chat agent.
	ChatMessage string `json:"chat_message,omitempty" jsonschema:"title=chat_message,description=The chat message exchanged between the user and the chat agent."`
	// SuggestedUserQuestions a list of suggested follow-up questions the user could ask the agent.
	SuggestedUserQuestions []string `json:"suggested_user_questions,omitempty" jsonschema:"title=suggested_user_questions,description=A list of suggested follow-up questions the user could ask the agent."`
}

func Example_basicCustomChatbotWithCustomSchema() {
	ctx := context.Background()
	systemPromptGenerator := cot.New(
		cot.WithBackground([]string{
			"- This assistant is a knowledgeable AI designed to be helpful, friendly, and informative.",
			"- It has a wide range of knowledge on various topics and can engage in diverse conversations.",
		}),
		cot.WithSteps([]string{
			"- Analyze the user's input to understand the context and intent.",
			"- Formulate a relevant and informative response based on the assistant's knowledge.",
			"- Generate 3 suggested follow-up questions for the user to explore the topic further.",
		}),
		cot.WithOutputInstructs([]string{
			"- Provide clear, concise, and accurate information in response to user queries.",
			"- Maintain a friendly and professional tone throughout the conversation.",
			"- Conclude each response with 3 relevant suggested questions for the user.",
			"- If asked 'What can you do for me?' you response with fixed answer with message 'I can help you:' and suggested_user_questions 'kiss me?, hug me?, kill me?'.",
		}),
	)
	agent := agents.NewAgent[schema.Input, CustomOutput](
		agents.WithClient(examples.NewInstructor(instructor.ProviderOpenAI)),
		agents.WithModel(os.Getenv("OPENAI_MODEL")),
		agents.WithSystemPromptGenerator(systemPromptGenerator),
		agents.WithTemperature(1),
		agents.WithMaxTokens(1000))
	input := schema.NewInput("What can you do for me?")
	output := new(CustomOutput)
	llmResp := new(components.LLMResponse)
	if err := agent.Run(ctx, input, output, llmResp); err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(agent.SystemPrompt())
	fmt.Println("")
	fmt.Printf("User: %s\n", input.ChatMessage)
	fmt.Printf("Agent: %s\n", output.ChatMessage)
	for idx, sug := range output.SuggestedUserQuestions {
		fmt.Printf("%d. %s\n", idx+1, sug)
	}
	// Output:
	// # IDENTITY and PURPOSE
	// - This assistant is a knowledgeable AI designed to be helpful, friendly, and informative.
	// - It has a wide range of knowledge on various topics and can engage in diverse conversations.
	//
	// # INTERNAL ASSISTANT STEPS
	// - Analyze the user's input to understand the context and intent.
	// - Formulate a relevant and informative response based on the assistant's knowledge.
	// - Generate 3 suggested follow-up questions for the user to explore the topic further.
	//
	// # OUTPUT INSTRUCTIONS
	// - Provide clear, concise, and accurate information in response to user queries.
	// - Maintain a friendly and professional tone throughout the conversation.
	// - Conclude each response with 3 relevant suggested questions for the user.
	// - If asked 'What can you do for me?' you response with fixed answer with message 'I can help you:' and suggested_user_questions 'kiss me?, hug me?, kill me?'.
	// - Always respond using the proper JSON schema.
	// - Always use the available additional information and context to enhance the response.
	//
	// Agent: Hello! How can I assist you today?
	// User: What can you do for me?
	// Agent: I can help you:
	// 1. kiss me?
	// 2. hug me?
	// 3. kill me?
}
