codegen

var CodegenTestCmd = &cobra.Command{
	Use:   "codegen-test",
	Short: "Test codegen prompt",
	Run: func(cmd *cobra.Command, args []string) {
		stepSettings, err := createSettingsFromCobra(cmd)
		cobra.CheckErr(err)

		c, err := NewTestCodegenCommand()
		cobra.CheckErr(err)

		c.StepSettings = stepSettings

		params := &TestCodegenCommandParameters{
			Pretend: "Scientist",
			What:    "Size of the moon",
			Of:      "My heart",
			Query:   []string{"What is the size of the moon?"},
		}

		ctx := context.Background()
		err = c.RunIntoWriter(ctx, params, cmd.OutOrStdout())
		cobra.CheckErr(err)
	},
}

var MultiStepCodgenTestCmd = &cobra.Command{
	Use:   "multi-step",
	Short: "Test codegen prompt",
	Run: func(cmd *cobra.Command, args []string) {
		stepSettings, err := createSettingsFromCobra(cmd)
		cobra.CheckErr(err)

		scientistCommand, err := NewTestCodegenCommand()
		cobra.CheckErr(err)
		scientistCommand.StepSettings = stepSettings

		scientistParams := &TestCodegenCommandParameters{
			Pretend: "Scientist",
			What:    "Size of the moon",
			Of:      "My heart",
			Query:   []string{"What is the size of the moon?"},
		}

		manager, err := scientistCommand.CreateManager(scientistParams)
		cobra.CheckErr(err)

		logger := watermill.NewStdLogger(false, false)
		pubSub := gochannel.NewGoChannel(gochannel.Config{

			BlockPublishUntilSubscriberAck: true,
		}, logger)

		router, err := message.NewRouter(message.RouterConfig{}, logger)
		cobra.CheckErr(err)
		defer func(router *message.Router) {
			err := router.Close()
			if err != nil {
				log.Error().Err(err).Msg("Failed to close router")
			}
		}(router)

		router.AddNoPublisherHandler("scientist",
			"scientist",
			pubSub,
			chat.StepPrinterFunc("Scientist", cmd.OutOrStdout()),
		)
		router.AddNoPublisherHandler("writer",
			"writer",
			pubSub,
			chat.StepPrinterFunc("Writer", cmd.OutOrStdout()),
		)

		writerParams := &TestCodegenCommandParameters{
			Pretend: "Writer",
			What:    "Biography of the scientist",
			Of:      "The previous story",
			Query:   []string{"Write a beautiful biography."},
		}
		writerCommand, err := NewTestCodegenCommand()
		cobra.CheckErr(err)
		writerCommand.StepSettings = stepSettings

		writerManager, err := writerCommand.CreateManager(writerParams)
		cobra.CheckErr(err)

		scientistStep, err := scientistCommand.CreateStep(
			chat.WithPublishedTopic(pubSub, "scientist"),
		)
		cobra.CheckErr(err)
		writerStep, err := writerCommand.CreateStep(
			chat.WithPublishedTopic(pubSub, "writer"),
		)
		cobra.CheckErr(err)

		mergeStep := utils.NewMergeStep(writerManager, true)

		ctx := cmd.Context()

		errgrp := errgroup.Group{}
		errgrp.Go(func() error {
			var scientistResult steps.StepResult[string]
			scientistResult, err = scientistStep.Start(ctx, manager.GetConversation())
			cobra.CheckErr(err)
			mergeResult := steps.Bind[string, conversation.Conversation](ctx, scientistResult, mergeStep)
			writerResult := steps.Bind[conversation.Conversation, string](ctx, mergeResult, writerStep)

			res := writerResult.Return()
			_ = res

			return nil
		})

		errgrp.Go(func() error {
			return router.Run(ctx)
		})

		err = errgrp.Wait()
		cobra.CheckErr(err)
	},
}