main.go

package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/sashabaranov/go-openai"
	"github.com/urfave/cli/v2"
)

const (
	Out    = "out"
	Prompt = "prompt"
)

const (
	envGogen      = "GOGEN_KEY"
	topDivider    = "<CODE>"
	bottomDivider = "</CODE>"
)

var system = `
You are golang developer in FAANG, Facebook, Apple, Amazon, Netflix, Google, Twitch.
You have in your context every development pattern, algorithm and data structures.
You are helping your colleague developers to write their code.

Please, answer only with golang code for every request.
You must not add any code, additional text or markdown declarations that user directly don't ask you to write

Usage examples must be written in golang comment blocks.

Code must be wrapped with tags <CODE> in the beginning and </CODE> in the end of your generated code. Example: <CODE>func generated() {}</CODE>
`

func main() {
	app := &cli.App{
		Name:        "gogen",
		Description: fmt.Sprintf("Environment variable %s is required to use OpenAPI api.", envGogen),
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    Out,
				Aliases: []string{"o"},
				Value:   fmt.Sprintf("gogen_%s.go", strconv.Itoa(int(time.Now().Unix()))),
				Usage:   "out file name",
			},
			&cli.StringFlag{
				Name:     Prompt,
				Aliases:  []string{"p"},
				Required: true,
				Usage:    "prompt text",
			},
		},
		Usage:  "use chatgpt to generate your code",
		Action: App,
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func App(c *cli.Context) error {
	k := os.Getenv(envGogen)
	if k == "" {
		return fmt.Errorf("%s is required. See --help for details", envGogen)
	}

	client := openai.NewClient(k)
	systemRequest := strings.Join(strings.Split(system, "\n"), " ")
	req := openai.ChatCompletionRequest{
		Model: openai.GPT4Turbo,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleSystem,
				Content: systemRequest,
			},
			{
				Role:    openai.ChatMessageRoleUser,
				Content: c.String(Prompt),
			},
		},
	}

	resp, err := client.CreateChatCompletion(c.Context, req)
	if err != nil {
		return fmt.Errorf("Completion error: %v\n", err)
	}

	f := strings.Builder{}
	f.WriteString(fmt.Sprintf("// Code generated by \"gogen %s\";\n", strings.Join(os.Args[1:], " ")))
	code, err := extractTheCode(resp.Choices[0].Message.Content)
	if err != nil {
		return fmt.Errorf("cannot extract code: %s", err)
	}

	if !strings.HasPrefix(code, "package") && !strings.HasPrefix(code, "\npackage") {
		f.WriteString(fmt.Sprintf("package gogen\n"))
	}

	f.WriteString(code)

	if err := os.WriteFile(c.String(Out), []byte(f.String()), 0644); err != nil {
		return fmt.Errorf("writing output: %s", err)
	}

	return nil
}

func extractTheCode(response string) (string, error) {
	if !strings.Contains(response, topDivider) || !strings.Contains(response, bottomDivider) {
		return "", errors.New(fmt.Sprintf("bad response: %s", response))
	}

	firstParts := strings.Split(response, topDivider)
	lastParts := strings.Split(firstParts[1], bottomDivider)
	return lastParts[0], nil
}