at-container-registry/pkg/atproto/generate.go

//go:build ignore
// +build ignore

package main

// Lexicon and CBOR Code Generator
//
// This generates:
// 1. Go types from lexicon JSON files (via lex/lexgen library)
// 2. CBOR marshaling code for ATProto records (via cbor-gen)
// 3. Type registration for lexutil (register.go)
//
// Usage:
//   go generate ./pkg/atproto/...
//
// Key insight: We use RegisterLexiconTypeID: false to avoid generating init()
// blocks that require CBORMarshaler. This breaks the circular dependency between
// lexgen and cbor-gen. See: https://github.com/bluesky-social/indigo/issues/931

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/bluesky-social/indigo/atproto/lexicon"
	"github.com/bluesky-social/indigo/lex/lexgen"
	"golang.org/x/tools/imports"
)

func main() {
	// Find repo root
	repoRoot, err := findRepoRoot()
	if err != nil {
		fmt.Printf("failed to find repo root: %v\n", err)
		os.Exit(1)
	}

	pkgDir := filepath.Join(repoRoot, "pkg/atproto")
	lexDir := filepath.Join(repoRoot, "lexicons")

	// Step 0: Clean up old register.go to avoid conflicts
	// (It will be regenerated at the end)
	os.Remove(filepath.Join(pkgDir, "register.go"))

	// Step 1: Load all lexicon schemas into catalog (for cross-references)
	fmt.Println("Loading lexicons...")
	cat := lexicon.NewBaseCatalog()
	if err := cat.LoadDirectory(lexDir); err != nil {
		fmt.Printf("failed to load lexicons: %v\n", err)
		os.Exit(1)
	}

	// Step 2: Generate Go code for each lexicon file
	fmt.Println("Running lexgen...")
	config := &lexgen.GenConfig{
		RegisterLexiconTypeID: false, // KEY: no init() blocks generated
		UnknownType:           "map-string-any",
		WarningText:           "Code generated by generate.go; DO NOT EDIT.",
	}

	// Track generated types for register.go
	var registeredTypes []typeInfo

	// Walk lexicon directory and generate code for each file
	err = filepath.Walk(lexDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() || !strings.HasSuffix(path, ".json") {
			return nil
		}

		// Load and parse the schema file
		data, err := os.ReadFile(path)
		if err != nil {
			return fmt.Errorf("failed to read %s: %w", path, err)
		}

		var sf lexicon.SchemaFile
		if err := json.Unmarshal(data, &sf); err != nil {
			return fmt.Errorf("failed to parse %s: %w", path, err)
		}

		if err := sf.FinishParse(); err != nil {
			return fmt.Errorf("failed to finish parse %s: %w", path, err)
		}

		// Flatten the schema
		flat, err := lexgen.FlattenSchemaFile(&sf)
		if err != nil {
			return fmt.Errorf("failed to flatten schema %s: %w", path, err)
		}

		// Generate code
		var buf bytes.Buffer
		gen := &lexgen.CodeGenerator{
			Config: config,
			Lex:    flat,
			Cat:    &cat,
			Out:    &buf,
		}

		if err := gen.WriteLexicon(); err != nil {
			return fmt.Errorf("failed to generate code for %s: %w", path, err)
		}

		// Fix package name: lexgen generates "ioatcr" but we want "atproto"
		code := bytes.Replace(buf.Bytes(), []byte("package ioatcr"), []byte("package atproto"), 1)

		// Format with goimports
		fileName := gen.FileName()
		formatted, err := imports.Process(fileName, code, nil)
		if err != nil {
			// Write unformatted for debugging
			outPath := filepath.Join(pkgDir, fileName)
			os.WriteFile(outPath+".broken", code, 0644)
			return fmt.Errorf("failed to format %s: %w (wrote to %s.broken)", fileName, err, outPath)
		}

		// Write output file
		outPath := filepath.Join(pkgDir, fileName)
		if err := os.WriteFile(outPath, formatted, 0644); err != nil {
			return fmt.Errorf("failed to write %s: %w", outPath, err)
		}

		fmt.Printf("  Generated %s\n", fileName)

		// Track type for registration - compute type name from NSID
		typeName := nsidToTypeName(sf.ID)
		registeredTypes = append(registeredTypes, typeInfo{
			NSID:     sf.ID,
			TypeName: typeName,
		})

		return nil
	})
	if err != nil {
		fmt.Printf("lexgen failed: %v\n", err)
		os.Exit(1)
	}

	// Step 3: Run cbor-gen via exec.Command
	// This must be a separate process so it can compile the freshly generated types
	fmt.Println("Running cbor-gen...")
	if err := runCborGen(repoRoot, pkgDir); err != nil {
		fmt.Printf("cbor-gen failed: %v\n", err)
		os.Exit(1)
	}

	// Step 4: Generate register.go
	fmt.Println("Generating register.go...")
	if err := generateRegisterFile(pkgDir, registeredTypes); err != nil {
		fmt.Printf("failed to generate register.go: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("Code generation complete!")
}

type typeInfo struct {
	NSID     string
	TypeName string
}

// nsidToTypeName converts an NSID to a Go type name
// io.atcr.manifest → Manifest
// io.atcr.hold.captain → HoldCaptain
// io.atcr.sailor.profile → SailorProfile
func nsidToTypeName(nsid string) string {
	parts := strings.Split(nsid, ".")
	if len(parts) < 3 {
		return ""
	}
	// Skip the first two parts (authority, e.g., "io.atcr")
	// and capitalize each remaining part
	var result string
	for _, part := range parts[2:] {
		if len(part) > 0 {
			result += strings.ToUpper(part[:1]) + part[1:]
		}
	}
	return result
}

func runCborGen(repoRoot, pkgDir string) error {
	// Create a temporary Go file that runs cbor-gen
	cborGenCode := `//go:build ignore

package main

import (
	"fmt"
	"os"

	cbg "github.com/whyrusleeping/cbor-gen"

	"atcr.io/pkg/atproto"
)

func main() {
	if err := cbg.WriteMapEncodersToFile("cbor_gen.go", "atproto",
		// Manifest types
		atproto.Manifest{},
		atproto.Manifest_BlobReference{},
		atproto.Manifest_ManifestReference{},
		atproto.Manifest_Platform{},
		atproto.Manifest_Annotations{},
		atproto.Manifest_BlobReference_Annotations{},
		atproto.Manifest_ManifestReference_Annotations{},
		// Tag
		atproto.Tag{},
		// Sailor types
		atproto.SailorProfile{},
		atproto.SailorStar{},
		atproto.SailorStar_Subject{},
		// Hold types
		atproto.HoldCaptain{},
		atproto.HoldCrew{},
		atproto.HoldLayer{},
		// External types
		atproto.TangledProfileRecord{},
	); err != nil {
		fmt.Printf("cbor-gen failed: %v\n", err)
		os.Exit(1)
	}
}
`

	// Write temp file
	tmpFile := filepath.Join(pkgDir, "cborgen_tmp.go")
	if err := os.WriteFile(tmpFile, []byte(cborGenCode), 0644); err != nil {
		return fmt.Errorf("failed to write temp cbor-gen file: %w", err)
	}
	defer os.Remove(tmpFile)

	// Run it
	cmd := exec.Command("go", "run", tmpFile)
	cmd.Dir = pkgDir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func generateRegisterFile(pkgDir string, types []typeInfo) error {
	var buf bytes.Buffer

	buf.WriteString("// Code generated by generate.go; DO NOT EDIT.\n\n")
	buf.WriteString("package atproto\n\n")
	buf.WriteString("import lexutil \"github.com/bluesky-social/indigo/lex/util\"\n\n")
	buf.WriteString("func init() {\n")

	for _, t := range types {
		fmt.Fprintf(&buf, "\tlexutil.RegisterType(%q, &%s{})\n", t.NSID, t.TypeName)
	}

	buf.WriteString("}\n")

	outPath := filepath.Join(pkgDir, "register.go")
	return os.WriteFile(outPath, buf.Bytes(), 0644)
}

func findRepoRoot() (string, error) {
	dir, err := os.Getwd()
	if err != nil {
		return "", err
	}

	for {
		if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
			return dir, nil
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			return "", fmt.Errorf("go.mod not found")
		}
		dir = parent
	}
}