package cmd

import (
	"fmt"

	cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util"
	"git.mlow.ca/mlow/lmcli/pkg/lmcli"
	"git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
	"github.com/spf13/cobra"
)

func EditCmd(ctx *lmcli.Context) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "edit <conversation>",
		Short: "Edit the last user reply in a conversation",
		Args: func(cmd *cobra.Command, args []string) error {
			argCount := 1
			if err := cobra.MinimumNArgs(argCount)(cmd, args); err != nil {
				return err
			}
			return nil
		},
		RunE: func(cmd *cobra.Command, args []string) error {
			shortName := args[0]
			conversation := cmdutil.LookupConversation(ctx, shortName)

			messages, err := ctx.Store.Messages(conversation)
			if err != nil {
				return fmt.Errorf("Could not retrieve messages for conversation: %s", conversation.Title)
			}

			offset, _ := cmd.Flags().GetInt("offset")
			if offset < 0 {
				offset = -offset
			}

			if offset > len(messages)-1 {
				return fmt.Errorf("Offset %d is before the start of the conversation.", offset)
			}

			desiredIdx := len(messages) - 1 - offset

			// walk backwards through the conversation deleting messages until and
			// including the last user message
			toRemove := []model.Message{}
			var toEdit *model.Message
			for i := len(messages) - 1; i >= 0; i-- {
				if i == desiredIdx {
					toEdit = &messages[i]
				}
				toRemove = append(toRemove, messages[i])
				messages = messages[:i]
				if toEdit != nil {
					break
				}
			}

			newContents := inputFromArgsOrEditor(args[1:], "# Save when finished editing\n", toEdit.Content)
			switch newContents {
			case toEdit.Content:
				return fmt.Errorf("No edits were made.")
			case "":
				return fmt.Errorf("No message was provided.")
			}

			role, _ := cmd.Flags().GetString("role")
			if role == "" {
				role = string(toEdit.Role)
			} else if role != string(model.MessageRoleUser) && role != string(model.MessageRoleAssistant) {
				return fmt.Errorf("Invalid role specified. Please use 'user' or 'assistant'.")
			}

			for _, message := range toRemove {
				err = ctx.Store.DeleteMessage(&message)
				if err != nil {
					lmcli.Warn("Could not delete message: %v\n", err)
				}
			}

			cmdutil.HandleConversationReply(ctx, conversation, true, model.Message{
				ConversationID:  conversation.ID,
				Role:            model.MessageRole(role),
				Content: newContents,
			})
			return nil
		},
		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
			compMode := cobra.ShellCompDirectiveNoFileComp
			if len(args) != 0 {
				return nil, compMode
			}
			return ctx.Store.ConversationShortNameCompletions(toComplete), compMode
		},
	}

	cmd.Flags().Int("offset", 1, "Offset from the last message to edit")
	cmd.Flags().StringP("role", "r", "", "Role of the edited message (user or assistant)")

	return cmd
}