ombi-mcp/tools/schemagen/main.go

// Command schemagen extracts the finalized JSON Schema contracts from
// docs/schema into internal/tools/schemas.json, which is embedded in
// the binary so tools/list can advertise complete input and output
// schemas without parsing markdown at runtime.
//
// Run from the repository root via `go generate ./internal/tools`.
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
)

const (
	inputDoc  = "docs/schema/03-input-schemas.md"
	defsDoc   = "docs/schema/04-results.md"
	outputDoc = "internal/tools/schemas.json"
)

var sectionRe = regexp.MustCompile(`^## ([a-z_]+)\s*$`)

// firstJSONBlock returns the first ```json fenced block following a
// line index, or nil if the section has none.
func firstJSONBlock(lines []string, start, end int) []byte {
	inBlock := false
	var buf bytes.Buffer
	for i := start; i < end; i++ {
		l := lines[i]
		switch {
		case !inBlock && l == "```json":
			inBlock = true
		case inBlock && l == "```":
			return buf.Bytes()
		case inBlock:
			buf.WriteString(l)
			buf.WriteByte('\n')
		}
	}
	return nil
}

func splitLines(b []byte) []string {
	var out []string
	start := 0
	for i := 0; i < len(b); i++ {
		if b[i] == '\n' {
			out = append(out, string(b[start:i]))
			start = i + 1
		}
	}
	if start < len(b) {
		out = append(out, string(b[start:]))
	}
	return out
}

func mustJSON(raw []byte, ctx string) json.RawMessage {
	var v any
	if err := json.Unmarshal(raw, &v); err != nil {
		fatal("%s: invalid JSON: %v", ctx, err)
	}
	return json.RawMessage(raw)
}

func fatal(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "schemagen: "+format+"\n", args...)
	os.Exit(1)
}

func main() {
	root, err := os.Getwd()
	if err != nil {
		fatal("cwd: %v", err)
	}

	// --- input schemas: one ```json block per `## tool` section ---
	inSrc, err := os.ReadFile(filepath.Join(root, inputDoc))
	if err != nil {
		fatal("read %s: %v", inputDoc, err)
	}
	lines := splitLines(inSrc)
	input := map[string]json.RawMessage{}
	for i, l := range lines {
		m := sectionRe.FindStringSubmatch(l)
		if m == nil {
			continue
		}
		name := m[1]
		end := len(lines)
		for j := i + 1; j < len(lines); j++ {
			if sectionRe.MatchString(lines[j]) || len(lines[j]) > 1 && lines[j][0] == '#' {
				end = j
				break
			}
		}
		block := firstJSONBlock(lines, i+1, end)
		if block == nil {
			fatal("%s: no json schema block for section %q", inputDoc, name)
		}
		input[name] = mustJSON(block, inputDoc+"/"+name)
	}
	if len(input) != 31 {
		fatal("expected 31 tool schemas, extracted %d", len(input))
	}

	// --- shared $defs library: first ```json block of 04-results.md ---
	defSrc, err := os.ReadFile(filepath.Join(root, defsDoc))
	if err != nil {
		fatal("read %s: %v", defsDoc, err)
	}
	dlines := splitLines(defSrc)
	defBlock := firstJSONBlock(dlines, 0, len(dlines))
	if defBlock == nil {
		fatal("%s: no json block found", defsDoc)
	}
	var defDoc struct {
		Defs map[string]json.RawMessage `json:"$defs"`
	}
	if err := json.Unmarshal(defBlock, &defDoc); err != nil {
		fatal("%s: invalid defs JSON: %v", defsDoc, err)
	}
	if len(defDoc.Defs) == 0 {
		fatal("%s: empty $defs", defsDoc)
	}

	out := struct {
		Input map[string]json.RawMessage `json:"input"`
		Defs  map[string]json.RawMessage `json:"defs"`
	}{Input: input, Defs: defDoc.Defs}
	raw, err := json.Marshal(out)
	if err != nil {
		fatal("marshal: %v", err)
	}
	var pretty bytes.Buffer
	if err := json.Indent(&pretty, raw, "", "  "); err != nil {
		fatal("indent: %v", err)
	}
	path := filepath.Join(root, outputDoc)
	if err := os.WriteFile(path, append(pretty.Bytes(), '\n'), 0o644); err != nil {
		fatal("write %s: %v", outputDoc, err)
	}
	names := make([]string, 0, len(input))
	for n := range input {
		names = append(names, n)
	}
	sort.Strings(names)
	fmt.Printf("wrote %s: %d input schemas, %d defs\n", outputDoc, len(input), len(defDoc.Defs))
}