package agents

import (
	"strings"

	gollm "gitea.stevedudenhoeffer.com/steve/go-llm"
)

// TidBit is a small piece of information that the AI has learned.
type TidBit struct {
	Info   string `json:"info"`
	Source string `json:"source"`
}

type Knowledge struct {
	// OriginalQuestions are the questions that was asked first to the AI before any processing was done.
	OriginalQuestions []string `json:"originalQuestions"`

	// RemainingQuestions is the questions that are left to find answers for.
	RemainingQuestions []string `json:"remainingQuestions"`

	// NotesToSelf are notes that the AI has made for itself.
	NotesToSelf []string `json:"notesToSelf"`

	// CurrentObjectives are the objectives that the AI is currently working on.
	CurrentObjectives []string `json:"currentObjectives"`

	// Knowledge are the tidbits of information that the AI has learned.
	Knowledge []TidBit `json:"knowledge"`
}

func (k Knowledge) ToSystemMessage() gollm.Message {
	var sources = map[string][]string{}

	for _, t := range k.Knowledge {
		sources[t.Source] = append(sources[t.Source], t.Info)
	}

	var msg string

	if len(k.OriginalQuestions) > 0 {
		msg += "Original questions asked:\n    - " + strings.Join(k.OriginalQuestions, "\n    - ") + "\n"
	}

	if len(k.NotesToSelf) > 0 {
		msg += "Notes to self:\n    - " + strings.Join(k.NotesToSelf, "\n    - ") + "\n"
	}

	if len(sources) > 0 {
		msg += "Learned information:\n"
		for source, info := range sources {
			if source == "" {
				source = "(unsourced)"
			}
			msg += "  - From " + source + ":\n    - " + strings.Join(info, "\n    - ") + "\n"
		}
	}

	if len(k.CurrentObjectives) > 0 {
		msg += "Current objectives:\n    - " + strings.Join(k.CurrentObjectives, "\n    - ") + "\n"
	}

	if len(k.RemainingQuestions) > 0 {
		msg += "Remaining questions:\n    - " + strings.Join(k.RemainingQuestions, "\n    - ") + "\n"
	}

	return gollm.Message{
		Role: gollm.RoleSystem,
		Text: msg,
	}
}

// ToMessage converts the knowledge to a message that can be sent to the LLM.
func (k Knowledge) ToMessage() string {
	var learned []string
	for _, t := range k.Knowledge {
		learned = append(learned, t.Info)
	}
	return "Original questions asked:\n" + strings.Join(k.OriginalQuestions, "\n") + "\n" +
		"Learned information:\n" + strings.Join(learned, "\n") + "\n" +
		"Remaining questions:\n" + strings.Join(k.RemainingQuestions, "\n")
}

// Absorb adds the knowledge from the other knowledge objects to this one.
// The OriginalQuestions returned will be from value of k's OriginalQuestions.
// If any absorbed knowledge has any "CurrentObjectives" set they will explicitly overwrite the current objectives in
// aggregate, all CurrentObjectives from all absorbed knowledge will be used.
// Any new Knowledge or NotesToSelf will be appended to the current Knowledge and NotesToSelf.
func (k Knowledge) Absorb(o ...Knowledge) Knowledge {
	res := k

	var newObjectives []string

	for _, ok := range o {
		if len(ok.CurrentObjectives) > 0 {
			newObjectives = append(newObjectives, ok.CurrentObjectives...)
		}

		res.NotesToSelf = append(res.NotesToSelf, ok.NotesToSelf...)
		res.Knowledge = append(res.Knowledge, ok.Knowledge...)
	}

	if len(newObjectives) > 0 {
		res.CurrentObjectives = newObjectives
	}

	return res
}