// Command schemagen generates Go types from the FLUID wire-contract schemas.
//
// The record types are never hand-written: schemas/ is the source of truth, and
// drift between the specification and the implementation must fail the build
// rather than survive as a quietly diverging struct.
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go/format"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// schema is the subset of JSON Schema the FLUID contract actually uses.
type schema struct {
	ID                   string             `json:"$id"`
	Ref                  string             `json:"$ref"`
	Title                string             `json:"title"`
	Description          string             `json:"description"`
	Type                 any                `json:"type"`
	Const                any                `json:"const"`
	Enum                 []string           `json:"enum"`
	Format               string             `json:"format"`
	Properties           map[string]*schema `json:"properties"`
	Required             []string           `json:"required"`
	Items                *schema            `json:"items"`
	Defs                 map[string]*schema `json:"$defs"`
	OneOf                []*schema          `json:"oneOf"`
	AdditionalProperties json.RawMessage    `json:"additionalProperties"`
	PropertyNames        *schema            `json:"propertyNames"`
}

// additional reports how additionalProperties was written: as a schema
// (a map value type), as false, or absent.
func (s *schema) additional() (*schema, bool) {
	if len(s.AdditionalProperties) == 0 {
		return nil, false
	}
	var b bool
	if err := json.Unmarshal(s.AdditionalProperties, &b); err == nil {
		return nil, b // true means "anything"; false means closed
	}
	var sub schema
	if err := json.Unmarshal(s.AdditionalProperties, &sub); err == nil {
		return &sub, true
	}
	return nil, false
}

func (s *schema) typeIs(want string) bool {
	switch t := s.Type.(type) {
	case string:
		return t == want
	case []any:
		for _, v := range t {
			if v == want {
				return true
			}
		}
	}
	return false
}

func (s *schema) nullable() bool {
	if s.typeIs("null") {
		return true
	}
	for _, o := range s.OneOf {
		if o.typeIs("null") {
			return true
		}
	}
	return false
}

var initialisms = map[string]string{
	"id": "ID", "ids": "IDs", "url": "URL", "uri": "URI", "api": "API",
	"ms": "MS", "http": "HTTP", "json": "JSON", "yaml": "YAML", "sha": "SHA",
	"p95": "P95", "iei": "IEI", "sla": "SLA", "ai": "AI",
}

// splitCamel inserts a separator at lowercase-to-uppercase boundaries so that
// camelCase schema keys such as "cohortId" split into words the initialism map
// can then correct.
func splitCamel(s string) string {
	var b strings.Builder
	runes := []rune(s)
	for i, r := range runes {
		if i > 0 && r >= 'A' && r <= 'Z' && runes[i-1] >= 'a' && runes[i-1] <= 'z' {
			b.WriteRune('_')
		}
		b.WriteRune(r)
	}
	return b.String()
}

func goName(s string) string {
	s = splitCamel(s)
	parts := strings.FieldsFunc(s, func(r rune) bool { return r == '_' || r == '-' || r == '.' })
	var b strings.Builder
	for _, p := range parts {
		if up, ok := initialisms[strings.ToLower(p)]; ok {
			b.WriteString(up)
			continue
		}
		if p == "" {
			continue
		}
		b.WriteString(strings.ToUpper(p[:1]))
		b.WriteString(p[1:])
	}
	return b.String()
}

// generator accumulates named types while walking the schemas.
type generator struct {
	commonDefs map[string]string // $defs key -> Go type name
	decls      []string
	seen       map[string]bool
}

func (g *generator) emit(decl string) {
	g.decls = append(g.decls, decl)
}

func comment(indent, text string) string {
	text = strings.Join(strings.Fields(text), " ")
	if text == "" {
		return ""
	}
	var out []string
	line := indent + "//"
	for _, w := range strings.Fields(text) {
		if len(line)+1+len(w) > 84 && line != indent+"//" {
			out = append(out, line)
			line = indent + "//"
		}
		line += " " + w
	}
	out = append(out, line)
	return strings.Join(out, "\n") + "\n"
}

// resolveRef maps "common.schema.json#/$defs/foo" to its generated Go type.
func (g *generator) resolveRef(ref string) (string, bool) {
	const marker = "#/$defs/"
	i := strings.Index(ref, marker)
	if i < 0 {
		return "", false
	}
	key := ref[i+len(marker):]
	name, ok := g.commonDefs[key]
	return name, ok
}

// goType returns the Go type for s, generating nested named types as needed.
// parent seeds the name of any anonymous struct or enum encountered.
func (g *generator) goType(parent, field string, s *schema, required bool) string {
	if s.Ref != "" {
		if name, ok := g.resolveRef(s.Ref); ok {
			if !required {
				return "*" + name
			}
			return name
		}
		return "any"
	}

	nullable := s.nullable()

	// oneOf that is just "X or null" collapses to a pointer to X.
	if len(s.OneOf) > 0 {
		for _, o := range s.OneOf {
			if o.typeIs("null") {
				continue
			}
			inner := g.goType(parent, field, o, true)
			return "*" + strings.TrimPrefix(inner, "*")
		}
		return "any"
	}

	// A union of several concrete types (guardrail thresholds may be number,
	// boolean or string) has no faithful Go equivalent but `any`.
	if ts, ok := s.Type.([]any); ok {
		concrete := 0
		for _, t := range ts {
			if t != "null" {
				concrete++
			}
		}
		if concrete > 1 {
			return "any"
		}
	}

	name := parent + goName(field)

	switch {
	case len(s.Enum) > 0:
		g.emitEnum(name, s)
		if !required || nullable {
			return "*" + name
		}
		return name

	case s.typeIs("object"):
		if len(s.Properties) > 0 {
			g.emitStruct(name, s)
			if !required || nullable {
				return "*" + name
			}
			return name
		}
		if sub, ok := s.additional(); ok && sub != nil {
			return "map[string]" + g.goType(name, "Value", sub, true)
		}
		return "map[string]any"

	case s.typeIs("array"):
		if s.Items == nil {
			return "[]any"
		}
		return "[]" + strings.TrimPrefix(g.goType(name, "Item", s.Items, true), "*")

	case s.typeIs("string"):
		base := "string"
		if s.Format == "date-time" {
			base = "time.Time"
		}
		if nullable || (!required && base == "time.Time") {
			return "*" + base
		}
		return base

	case s.typeIs("integer"):
		if nullable || !required {
			return "*int64"
		}
		return "int64"

	case s.typeIs("number"):
		if nullable || !required {
			return "*float64"
		}
		return "float64"

	case s.typeIs("boolean"):
		if nullable || !required {
			return "*bool"
		}
		return "bool"
	}

	return "any"
}

// symbolNames covers enum values that are operators rather than words, so the
// generated constants stay valid Go identifiers.
var symbolNames = map[string]string{
	"<": "Lt", "<=": "Lte", "==": "Eq", "!=": "Neq", ">=": "Gte", ">": "Gt",
}

// enumConstName builds an identifier suffix for one enum value.
func enumConstName(v string, index int) string {
	if sym, ok := symbolNames[v]; ok {
		return sym
	}
	name := goName(v)
	cleaned := make([]rune, 0, len(name))
	for _, r := range name {
		if r == '_' || r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' {
			cleaned = append(cleaned, r)
		}
	}
	name = string(cleaned)
	if name == "" || name[0] >= '0' && name[0] <= '9' {
		name = fmt.Sprintf("Value%d", index)
	}
	return name
}

func (g *generator) emitEnum(name string, s *schema) {
	if g.seen[name] {
		return
	}
	g.seen[name] = true

	var b strings.Builder
	b.WriteString(comment("", s.Description))
	fmt.Fprintf(&b, "type %s string\n\nconst (\n", name)
	for i, v := range s.Enum {
		fmt.Fprintf(&b, "\t%s%s %s = %q\n", name, enumConstName(v, i), name, v)
	}
	b.WriteString(")\n\n")

	// A generated Valid method keeps enum checking in one place.
	fmt.Fprintf(&b, "// Valid reports whether v is a defined %s.\n", name)
	fmt.Fprintf(&b, "func (v %s) Valid() bool {\n\tswitch v {\n\tcase ", name)
	quoted := make([]string, 0, len(s.Enum))
	for i, v := range s.Enum {
		quoted = append(quoted, name+enumConstName(v, i))
	}
	b.WriteString(strings.Join(quoted, ", "))
	b.WriteString(":\n\t\treturn true\n\t}\n\treturn false\n}\n")

	g.emit(b.String())
}

func (g *generator) emitStruct(name string, s *schema) {
	if g.seen[name] {
		return
	}
	g.seen[name] = true

	req := map[string]bool{}
	for _, r := range s.Required {
		req[r] = true
	}

	keys := make([]string, 0, len(s.Properties))
	for k := range s.Properties {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	// Fields are resolved before the struct is written so nested types are
	// declared in dependency order.
	type field struct{ name, typ, tag, doc string }
	fields := make([]field, 0, len(keys))
	for _, k := range keys {
		p := s.Properties[k]
		typ := g.goType(name, k, p, req[k])
		tag := k
		if !req[k] {
			tag += ",omitempty"
		}
		fields = append(fields, field{goName(k), typ, tag, p.Description})
	}

	var b strings.Builder
	b.WriteString(comment("", s.Description))
	fmt.Fprintf(&b, "type %s struct {\n", name)
	for i, f := range fields {
		if f.doc != "" {
			if i > 0 {
				b.WriteString("\n")
			}
			b.WriteString(comment("\t", f.doc))
		}
		fmt.Fprintf(&b, "\t%s %s `json:%q yaml:%q`\n", f.name, f.typ, f.tag, f.tag)
	}
	b.WriteString("}\n")

	g.emit(b.String())
}

func load(path string) (*schema, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var s schema
	if err := json.Unmarshal(data, &s); err != nil {
		return nil, fmt.Errorf("%s: %w", filepath.Base(path), err)
	}
	return &s, nil
}

func write(path, pkg string, decls []string) error {
	var b bytes.Buffer
	b.WriteString("// Code generated by tools/schemagen. DO NOT EDIT.\n")
	b.WriteString("// Source: schemas/. Regenerate with `make generate`.\n\n")
	fmt.Fprintf(&b, "package %s\n\n", pkg)
	body := strings.Join(decls, "\n")
	if strings.Contains(body, "time.Time") {
		b.WriteString("import \"time\"\n\n")
	}
	b.WriteString(body)

	src, err := format.Source(b.Bytes())
	if err != nil {
		// Write the unformatted source so the error is diagnosable.
		_ = os.WriteFile(path+".broken", b.Bytes(), 0o644)
		return fmt.Errorf("format %s: %w", path, err)
	}
	return os.WriteFile(path, src, 0o644)
}

func main() {
	schemaDir := "schemas"
	outDir := filepath.Join("internal", "contract")

	if err := os.MkdirAll(outDir, 0o755); err != nil {
		fatal(err)
	}

	common, err := load(filepath.Join(schemaDir, "common.schema.json"))
	if err != nil {
		fatal(err)
	}

	g := &generator{commonDefs: map[string]string{}, seen: map[string]bool{}}

	// Pass one: name every common definition so record schemas can reference them.
	defKeys := make([]string, 0, len(common.Defs))
	for k := range common.Defs {
		defKeys = append(defKeys, k)
	}
	sort.Strings(defKeys)
	for _, k := range defKeys {
		g.commonDefs[k] = goName(k)
	}

	// Pass two: declare them.
	for _, k := range defKeys {
		d := common.Defs[k]
		name := g.commonDefs[k]
		switch {
		case len(d.Enum) > 0:
			g.emitEnum(name, d)
		case d.typeIs("object") && len(d.Properties) > 0:
			g.emitStruct(name, d)
		default:
			if g.seen[name] {
				continue
			}
			g.seen[name] = true
			underlying := g.goType("", k, d, true)
			// Scalar definitions become distinct types rather than aliases:
			// passing a hypothesis id where a revision id belongs should not
			// compile. Everything else stays an alias to keep call sites plain.
			form := "type %s = %s\n"
			switch underlying {
			case "string", "float64", "int64":
				form = "type %s %s\n"
			}
			g.emit(comment("", d.Description) + fmt.Sprintf(form, name, underlying))
		}
	}

	if err := write(filepath.Join(outDir, "common_gen.go"), "contract", g.decls); err != nil {
		fatal(err)
	}
	fmt.Printf("generated %s (%d declarations)\n", filepath.Join(outDir, "common_gen.go"), len(g.decls))

	records := []string{
		"pressure", "hypothesis", "revision", "experiment", "event",
		"feedback", "backend-requirement", "revision-descriptor",
		"routing-policy", "telemetry-envelope",
	}

	for _, rec := range records {
		s, err := load(filepath.Join(schemaDir, rec+".schema.json"))
		if err != nil {
			fatal(err)
		}
		rg := &generator{commonDefs: g.commonDefs, seen: map[string]bool{}}

		// Each record schema is a single-property wrapper; generate the
		// wrapper plus the payload type it names.
		for key, prop := range s.Properties {
			rg.emitStruct(goName(key), prop)
		}
		var keys []string
		for k := range s.Properties {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		var doc strings.Builder
		doc.WriteString(comment("", s.Description))
		fmt.Fprintf(&doc, "type %sDocument struct {\n", goName(rec))
		for _, k := range keys {
			fmt.Fprintf(&doc, "\t%s %s `json:%q yaml:%q`\n", goName(k), goName(k), k, k)
		}
		doc.WriteString("}\n")
		rg.emit(doc.String())

		out := filepath.Join(outDir, strings.ReplaceAll(rec, "-", "_")+"_gen.go")
		if err := write(out, "contract", rg.decls); err != nil {
			fatal(err)
		}
		fmt.Printf("generated %s\n", out)
	}
}

func fatal(err error) {
	fmt.Fprintln(os.Stderr, "schemagen:", err)
	os.Exit(1)
}