package console_new

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/docker/client"

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

	"gitea.stevedudenhoeffer.com/steve/answer/pkg/agents"
	"gitea.stevedudenhoeffer.com/steve/answer/pkg/agents/shared"
)

type Agent struct {
	agents.Agent
	// Model is the chat completion model to use
	Model gollm.ChatCompletion

	OnLoopComplete func(ctx context.Context, knowledge shared.Knowledge) error

	OnCommandStart func(ctx context.Context, command string) error
	OnCommandDone  func(ctx context.Context, command string, output string, err error) error

	OnDone func(ctx context.Context, knowledge shared.Knowledge) error

	ContextualInformation []string

	MaxCommands int
}

type Response struct {
	Knowledge shared.Knowledge
	DataDir   string
	OutputDir string
}

// Answer will give the model access to an ubuntu console with python and pip installed, and then ask the model to
// do what is necessary to answer the question.
func (a Agent) Answer(ctx context.Context, questions []string) (Response, error) {
	var res Response

	a.Agent = agents.NewAgent(a.Model, gollm.NewToolBox()).WithMaxCalls(200)

	if a.MaxCommands <= 0 {
		a.MaxCommands = 20 // Default to 20 commands as per requirements
	}

	res.Knowledge = shared.Knowledge{
		OriginalQuestions:  questions,
		RemainingQuestions: questions,
	}

	// create temporary directories for data and output
	dataDir, err := os.MkdirTemp("", "console-data-")
	if err != nil {
		return res, err
	}

	outputDir, err := os.MkdirTemp("", "console-output-")
	if err != nil {
		os.RemoveAll(dataDir)
		return res, err
	}

	res.DataDir = dataDir
	res.OutputDir = outputDir

	cl, err := client.NewClientWithOpts(client.FromEnv)
	if err != nil {
		os.RemoveAll(dataDir)
		os.RemoveAll(outputDir)
		return res, err
	}
	defer cl.Close()

	mounts := []mount.Mount{
		{
			Type:   mount.TypeBind,
			Source: dataDir,
			Target: "/data",
		},
		{
			Type:   mount.TypeBind,
			Source: outputDir,
			Target: "/output",
		},
	}

	c, err := CreateContainer(ctx, cl, ContainerConfig{
		Config: &container.Config{
			Image:      "ubuntu:latest",
			Cmd:        []string{"tail", "-f", "/dev/null"},
			Tty:        true,
			WorkingDir: "/data",
		},
		HostConfig: &container.HostConfig{
			AutoRemove: true,
			Mounts:     mounts,
		},
		Name: filepath.Base(dataDir),
	})

	if err != nil {
		os.RemoveAll(dataDir)
		os.RemoveAll(outputDir)
		return res, fmt.Errorf("failed to create container: %w", err)
	}
	defer func() {
		_ = c.Close(ctx)
	}()

	slog.Info("starting container", "dataDir", dataDir, "outputDir", outputDir, "container", fmt.Sprintf("%+v", c))
	err = c.Start(ctx)
	if err != nil {
		os.RemoveAll(dataDir)
		os.RemoveAll(outputDir)
		return res, fmt.Errorf("failed to start container: %w", err)
	}

	// Run the model

	var history executions
	var keepGoing = true

	// Initial setup - install basic tools
	setupCmd, setupErr := c.Sudo(ctx, "apt-get update && apt-get install -y curl wget git python3 python3-pip")
	if setupErr == nil {
		history = append(history, execution{
			Command:      "sudo apt-get update && apt-get install -y curl wget git python3 python3-pip",
			Output:       setupCmd,
			WhatILearned: []string{"Basic tools installed: curl, wget, git, python3, pip"},
		})
	}

	tools := map[string]gollm.Function{
		"exit": gollm.NewFunction(
			"exit",
			"exit the container",
			func(ctx *gollm.Context, args struct {
				RemainingQuestions []string `description:"any remaining questions that remain unanswered"`
			}) (any, error) {
				keepGoing = false
				return "exiting", nil
			}),

		"write_data": gollm.NewFunction(
			"write_data",
			"write a file in the /data directory",
			func(ctx *gollm.Context, args struct {
				Filename string `description:"The name of the file to write"`
				Content  string `description:"The content of the file to write"`
			}) (any, error) {
				target, err := SafeJoinPath(dataDir, args.Filename)
				if err != nil {
					return "", err
				}

				// Ensure directory exists
				dir := filepath.Dir(target)
				if err := os.MkdirAll(dir, 0755); err != nil {
					return "", err
				}

				f, err := os.Create(target)
				if err != nil {
					return "", err
				}
				defer f.Close()

				_, err = f.WriteString(args.Content)
				if err != nil {
					return "", err
				}

				return "wrote file to /data/" + args.Filename, nil
			}),

		"write_output": gollm.NewFunction(
			"write_output",
			"write a file in the /output directory (files here will be available after the agent completes)",
			func(ctx *gollm.Context, args struct {
				Filename string `description:"The name of the file to write"`
				Content  string `description:"The content of the file to write"`
			}) (any, error) {
				target, err := SafeJoinPath(outputDir, args.Filename)
				if err != nil {
					return "", err
				}

				// Ensure directory exists
				dir := filepath.Dir(target)
				if err := os.MkdirAll(dir, 0755); err != nil {
					return "", err
				}

				f, err := os.Create(target)
				if err != nil {
					return "", err
				}
				defer f.Close()

				_, err = f.WriteString(args.Content)
				if err != nil {
					return "", err
				}

				return "wrote file to /output/" + args.Filename, nil
			}),

		"read_data": gollm.NewFunction(
			"read_data",
			"read a file from the /data directory",
			func(ctx *gollm.Context, args struct {
				Filename string `description:"The name of the file to read"`
			}) (any, error) {
				target, err := SafeJoinPath(dataDir, args.Filename)
				if err != nil {
					return "", err
				}

				b, err := os.ReadFile(target)
				if err != nil {
					return "", err
				}

				return string(b), nil
			}),

		"read_output": gollm.NewFunction(
			"read_output",
			"read a file from the /output directory",
			func(ctx *gollm.Context, args struct {
				Filename string `description:"The name of the file to read"`
			}) (any, error) {
				target, err := SafeJoinPath(outputDir, args.Filename)
				if err != nil {
					return "", err
				}

				b, err := os.ReadFile(target)
				if err != nil {
					return "", err
				}

				return string(b), nil
			}),

		"execute": gollm.NewFunction(
			"execute",
			"execute a command in the container",
			func(ctx *gollm.Context, args struct {
				Command string   `description:"The command to execute"`
				Learn   []string `description:"What you learned from this command (optional)"`
				ToLearn []string `description:"What you still need to learn (optional)"`
			}) (any, error) {
				if len(history) >= a.MaxCommands {
					return "Command limit reached. You've used all " + fmt.Sprintf("%d", a.MaxCommands) + " allowed commands.", nil
				}

				if a.OnCommandStart != nil {
					err := a.OnCommandStart(ctx, args.Command)
					if err != nil {
						return "", err
					}
				}

				var res string
				// if the command starts with sudo then we need to use the sudo function
				if strings.HasPrefix(args.Command, "sudo ") {
					res, err = c.Sudo(ctx, args.Command[5:])
				} else {
					res, err = c.Execute(ctx, args.Command)
				}

				if a.OnCommandDone != nil {
					err = a.OnCommandDone(ctx, args.Command, res, nil)
					if err != nil {
						return "", err
					}
				}

				history = append(history, execution{
					Command:               args.Command,
					Output:                res,
					WhatILearned:          args.Learn,
					WhatIStillNeedToLearn: args.ToLearn,
				})

				return res, nil
			}),

		"summarize_knowledge": gollm.NewFunction(
			"summarize_knowledge",
			"summarize what you've learned so far",
			func(ctx *gollm.Context, args struct{}) (any, error) {
				var learned []string
				var toLearn []string

				for _, exec := range history {
					learned = append(learned, exec.WhatILearned...)
					toLearn = append(toLearn, exec.WhatIStillNeedToLearn...)
				}

				summary := "Knowledge Summary:\n"
				if len(learned) > 0 {
					summary += "What I've learned:\n- " + strings.Join(learned, "\n- ") + "\n\n"
				} else {
					summary += "I haven't learned anything specific yet.\n\n"
				}

				if len(toLearn) > 0 {
					summary += "What I still need to learn:\n- " + strings.Join(toLearn, "\n- ")
				} else {
					summary += "I don't have any specific learning goals at the moment."
				}

				return summary, nil
			}),
	}

	for i := 0; i < a.MaxCommands && len(history) < a.MaxCommands && keepGoing; i++ {

		systemPrompt := `You are now in a shell in a container of the ubuntu:latest image to answer a question asked by the user. 
You have full control over a bash shell inside this Docker container.

Important directories:
- /data: A temporary directory with the lifespan of your processing. Use this for working files.
- /output: Files placed here will be returned to the caller after you're done. Use this for final results.

You can execute up to ` + fmt.Sprintf("%d", a.MaxCommands) + ` commands total. You're currently on command ` + fmt.Sprintf("%d", len(history)+1) + ` of ` + fmt.Sprintf("%d", a.MaxCommands) + `.

For each command, you should:
1. Think about what you need to learn
2. Execute the command using the "execute" function
3. Analyze the output and record what you learned
4. Plan your next command based on this knowledge

You can write files directly to /data or /output using the write_data and write_output functions.
When you are done, use "exit" to exit the container.`

		var toolbox []gollm.Function

		// Add all tools
		toolbox = append(toolbox,
			tools["exit"],
			tools["write_data"],
			tools["write_output"],
			tools["read_data"],
			tools["read_output"],
			tools["summarize_knowledge"],
		)

		// Only add execute if we haven't reached the command limit
		if len(history) < a.MaxCommands {
			toolbox = append(toolbox, tools["execute"])
		}

		kw := shared.KnowledgeWorker{
			Model:                 a.Model,
			ToolBox:               gollm.NewToolBox(toolbox...),
			ContextualInformation: a.ContextualInformation,
			OnNewFunction: func(ctx context.Context, funcName string, args string) (any, error) {
				slog.Info("new function called", "function name", funcName, "args", args)
				return nil, nil
			},
		}

		slog.Info("answering question", "question", questions[0])
		r, err := kw.Answer(ctx, &res.Knowledge, systemPrompt, "", "", history.ToGeneralButLastMessageHistory(), func(res gollm.ToolCallResponse) {
		})

		if err != nil {
			return res, fmt.Errorf("error answering question: %w", err)
		}

		if len(r.Knowledge) > 0 {
			slog.Info("answered question and learned", "knowledge", r.Knowledge)
		} else {
			slog.Info("answered question and learned nothing")
		}

		res.Knowledge, err = a.KnowledgeIntegrate(ctx, res.Knowledge, r)
		if err != nil {
			return res, fmt.Errorf("error integrating knowledge: %w", err)
		}

		slog.Info("knowledge integrated", "question", questions[0], "knowledge", res.Knowledge)
	}
	return res, nil
}