output.go

package codegen

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"text/template"

	"github.com/oapi-codegen/oapi-codegen-exp/codegen/internal/templates"
	"golang.org/x/tools/imports"
)

// Output collects generated Go code and formats it.
type Output struct {
	packageName string
	imports     map[string]string // path -> alias
	types       []string          // type definitions in order
}

// NewOutput creates a new output collector.
func NewOutput(packageName string) *Output {
	return &Output{
		packageName: packageName,
		imports:     make(map[string]string),
	}
}

// AddImport adds an import path with optional alias.
func (o *Output) AddImport(path, alias string) {
	if path == "" {
		return
	}
	o.imports[path] = alias
}

// AddImports adds multiple imports from a map.
func (o *Output) AddImports(imports map[string]string) {
	for path, alias := range imports {
		o.AddImport(path, alias)
	}
}

// AddType adds a type definition to the output.
func (o *Output) AddType(code string) {
	if code != "" {
		o.types = append(o.types, code)
	}
}

// String generates the complete Go source file.
func (o *Output) String() string {
	var buf bytes.Buffer

	// Generated code header (tells linters to skip this file)
	buf.WriteString("// Code generated by oapi-codegen; DO NOT EDIT.\n\n")

	// Package declaration
	fmt.Fprintf(&buf, "package %s\n\n", o.packageName)

	// Imports
	if len(o.imports) > 0 {
		buf.WriteString("import (\n")
		paths := make([]string, 0, len(o.imports))
		for path := range o.imports {
			paths = append(paths, path)
		}
		sort.Strings(paths)

		for _, path := range paths {
			alias := o.imports[path]
			if alias != "" {
				fmt.Fprintf(&buf, "\t%s %q\n", alias, path)
			} else {
				fmt.Fprintf(&buf, "\t%q\n", path)
			}
		}
		buf.WriteString(")\n\n")
	}

	// Types
	for _, t := range o.types {
		buf.WriteString(t)
		buf.WriteString("\n\n")
	}

	return buf.String()
}

// Format returns the formatted Go source code with imports organized.
func (o *Output) Format() (string, error) {
	src := o.String()
	formatted, err := imports.Process("", []byte(src), nil)
	if err != nil {
		return src, fmt.Errorf("formatting output: %w (source:\n%s)", err, src)
	}
	return string(formatted), nil
}

// CodeBuilder helps construct Go code fragments.
type CodeBuilder struct {
	buf    bytes.Buffer
	indent int
}

// NewCodeBuilder creates a new code builder.
func NewCodeBuilder() *CodeBuilder {
	return &CodeBuilder{}
}

// Indent increases indentation.
func (b *CodeBuilder) Indent() {
	b.indent++
}

// Dedent decreases indentation.
func (b *CodeBuilder) Dedent() {
	if b.indent > 0 {
		b.indent--
	}
}

// Line writes a line with current indentation.
func (b *CodeBuilder) Line(format string, args ...any) {
	for i := 0; i < b.indent; i++ {
		b.buf.WriteByte('\t')
	}
	if len(args) > 0 {
		fmt.Fprintf(&b.buf, format, args...)
	} else {
		b.buf.WriteString(format)
	}
	b.buf.WriteByte('\n')
}

// BlankLine writes an empty line.
func (b *CodeBuilder) BlankLine() {
	b.buf.WriteByte('\n')
}

// Raw writes raw text without indentation or newline.
func (b *CodeBuilder) Raw(s string) {
	b.buf.WriteString(s)
}

// String returns the built code.
func (b *CodeBuilder) String() string {
	return b.buf.String()
}

// GenerateStruct generates a struct type definition.
func GenerateStruct(name string, fields []StructField, doc string, tagGen *StructTagGenerator) string {
	b := NewCodeBuilder()

	// Type documentation
	if doc != "" {
		for _, line := range strings.Split(doc, "\n") {
			b.Line("// %s", line)
		}
	}

	b.Line("type %s struct {", name)
	b.Indent()

	for _, f := range fields {
		tag := generateFieldTag(f, tagGen)
		if f.Doc != "" {
			for _, line := range strings.Split(f.Doc, "\n") {
				line = strings.TrimRight(line, " \t")
				b.Line("// %s", line)
			}
		}
		b.Line("%s %s %s", f.Name, f.Type, tag)
	}

	b.Dedent()
	b.Line("}")

	return b.String()
}

// generateFieldTag generates the struct tag for a field.
// All tags go through the template engine with a simple 2-field context.
// Extension-driven overrides (JSONIgnore, OmitZero, OmitEmpty) are applied
// as post-processing on the resulting tag map.
func generateFieldTag(f StructField, tagGen *StructTagGenerator) string {
	info := StructTagInfo{
		FieldName:  f.JSONName,
		IsOptional: !f.Required,
	}

	// All tags through the same template engine
	tags := tagGen.GenerateTagsMap(info)

	// Extension overrides on well-known tags
	if f.JSONIgnore {
		tags["json"] = "-"
		tags["form"] = "-"
	} else {
		// OmitEmpty extension override: if extensions explicitly set OmitEmpty
		// differently from the schema default (!f.Required), adjust tags.
		if f.OmitEmpty != !f.Required {
			applyOmitEmptyOverride(tags, f.JSONName, f.OmitEmpty, "json", "form")
		}
		// OmitZero (json-specific)
		if f.OmitZero {
			if v, ok := tags["json"]; ok {
				tags["json"] = v + ",omitzero"
			}
		}
	}

	return FormatTagsMap(tags)
}

// GenerateStructWithAdditionalProps generates a struct with AdditionalProperties field
// and custom marshal/unmarshal methods.
func GenerateStructWithAdditionalProps(name string, fields []StructField, addPropsType string, doc string, tagGen *StructTagGenerator) string {
	b := NewCodeBuilder()

	// Type documentation
	if doc != "" {
		for _, line := range strings.Split(doc, "\n") {
			b.Line("// %s", line)
		}
	}

	b.Line("type %s struct {", name)
	b.Indent()

	// Regular fields
	for _, f := range fields {
		tag := generateFieldTag(f, tagGen)
		b.Line("%s %s %s", f.Name, f.Type, tag)
	}

	// AdditionalProperties field
	b.Line("AdditionalProperties map[string]%s `json:\"-\"`", addPropsType)

	b.Dedent()
	b.Line("}")

	return b.String()
}

// GenerateTypeAlias generates a type alias definition.
func GenerateTypeAlias(name, targetType, doc string) string {
	b := NewCodeBuilder()

	if doc != "" {
		for _, line := range strings.Split(doc, "\n") {
			b.Line("// %s", line)
		}
	}

	b.Line("type %s = %s", name, targetType)

	return b.String()
}

// GenerateEnumFromInfo generates an enum type with const values using pre-computed EnumInfo.
// The EnumInfo contains sanitized names and the prefix decision from collision detection.
func GenerateEnumFromInfo(info *EnumInfo) string {
	b := NewCodeBuilder()

	if info.Doc != "" {
		for _, line := range strings.Split(info.Doc, "\n") {
			b.Line("// %s", line)
		}
	}

	b.Line("type %s %s", info.TypeName, info.BaseType)
	b.BlankLine()

	if len(info.Values) > 0 {
		b.Line("const (")
		b.Indent()

		for i, v := range info.Values {
			var constName string
			if info.PrefixTypeName {
				constName = info.TypeName + info.SanitizedNames[i]
			} else {
				constName = info.SanitizedNames[i]
			}

			if i < len(info.ValueDocs) && info.ValueDocs[i] != "" {
				for _, line := range strings.Split(info.ValueDocs[i], "\n") {
					b.Line("// %s", line)
				}
			}

			if info.BaseType == "string" {
				b.Line("%s %s = %q", constName, info.TypeName, v)
			} else {
				b.Line("%s %s = %s", constName, info.TypeName, v)
			}
		}

		b.Dedent()
		b.Line(")")
	}

	return b.String()
}

// UnionMember represents a member of a union type (anyOf/oneOf).
type UnionMember struct {
	TypeName            string   // Go type name (e.g., "Cat")
	MethodName          string   // Suffix for As/From/Merge methods (e.g., "Cat")
	Index               int      // Position in anyOf/oneOf array
	HasApplyDefaults    bool     // Whether this type has an ApplyDefaults method
	DiscriminatorValues []string // Discriminator mapping keys for this variant (empty if unmapped)
}

// UnionTypeConfig holds all information needed to generate a union type.
type UnionTypeConfig struct {
	TypeName      string
	Members       []UnionMember
	IsOneOf       bool
	Doc           string
	FixedFields   []StructField
	Discriminator *DiscriminatorInfo
	TagGen        *StructTagGenerator
	HelperPrefix  string         // e.g., "oapiCodegenHelpersPkg." or "" for embedded
	Converter     *NameConverter // for converting JSON property names to Go field names
}

// hasFixedField returns true if the given JSON field name is among the fixed fields.
func (cfg UnionTypeConfig) hasFixedField(jsonName string) bool {
	for _, f := range cfg.FixedFields {
		if f.JSONName == jsonName {
			return true
		}
	}
	return false
}

// unionTemplateFixedField is a pre-computed fixed field for the union template.
type unionTemplateFixedField struct {
	Name     string // Go field name
	Type     string // Go type
	JSONName string // JSON property name
	Tag      string // Full struct tag string
	Doc      string // Doc comment
	Required bool   // Whether the field is required
}

// unionTemplateMember is a pre-computed member for the union template.
type unionTemplateMember struct {
	TypeName             string // Go type name
	MethodName           string // Suffix for As/From/Merge
	DiscriminatorAutoSet string // Pre-computed auto-set line (e.g., `v.AuthType = "none"`) or empty
}

// unionTemplateDiscEntry is a discriminator value → method mapping for ValueByDiscriminator.
type unionTemplateDiscEntry struct {
	Value      string
	MethodName string
}

// unionTemplateData is the data passed to the union template.
type unionTemplateData struct {
	TypeName             string
	Doc                  string
	HelperPrefix         string
	FixedFields          []unionTemplateFixedField
	Members              []unionTemplateMember
	Discriminator        *DiscriminatorInfo
	DiscriminatorEntries []unionTemplateDiscEntry
}

// GenerateUnionCode generates all code for a union type using the union template.
func GenerateUnionCode(cfg UnionTypeConfig) (string, error) {
	data := buildUnionTemplateData(cfg)

	tmpl, err := loadUnionTemplate()
	if err != nil {
		return "", err
	}

	var buf strings.Builder

	// Type definition
	if err := tmpl.ExecuteTemplate(&buf, "union_type", data); err != nil {
		return "", fmt.Errorf("executing union_type: %w", err)
	}

	// Accessors (As/From/Merge)
	if err := tmpl.ExecuteTemplate(&buf, "union_accessors", data); err != nil {
		return "", fmt.Errorf("executing union_accessors: %w", err)
	}

	// Discriminator methods
	if data.Discriminator != nil {
		if err := tmpl.ExecuteTemplate(&buf, "union_discriminator", data); err != nil {
			return "", fmt.Errorf("executing union_discriminator: %w", err)
		}
	}

	// Marshal/Unmarshal
	if len(data.FixedFields) > 0 {
		if err := tmpl.ExecuteTemplate(&buf, "union_marshal_fixed_fields", data); err != nil {
			return "", fmt.Errorf("executing union_marshal_fixed_fields: %w", err)
		}
	} else {
		if err := tmpl.ExecuteTemplate(&buf, "union_marshal_simple", data); err != nil {
			return "", fmt.Errorf("executing union_marshal_simple: %w", err)
		}
	}

	// ApplyDefaults
	if err := tmpl.ExecuteTemplate(&buf, "union_apply_defaults", data); err != nil {
		return "", fmt.Errorf("executing union_apply_defaults: %w", err)
	}

	return buf.String(), nil
}

// buildUnionTemplateData pre-computes all values the template needs.
func buildUnionTemplateData(cfg UnionTypeConfig) unionTemplateData {
	// Pre-compute fixed fields with tags
	var fixedFields []unionTemplateFixedField
	for _, f := range cfg.FixedFields {
		fixedFields = append(fixedFields, unionTemplateFixedField{
			Name:     f.Name,
			Type:     f.Type,
			JSONName: f.JSONName,
			Tag:      generateFieldTag(f, cfg.TagGen),
			Doc:      f.Doc,
			Required: f.Required,
		})
	}

	// Pre-compute discriminator auto-set for each member
	allMapped := allMembersMapped(cfg)
	var members []unionTemplateMember
	for _, m := range cfg.Members {
		tm := unionTemplateMember{
			TypeName:   m.TypeName,
			MethodName: m.MethodName,
		}
		if allMapped && len(m.DiscriminatorValues) > 0 {
			tm.DiscriminatorAutoSet = computeDiscriminatorAutoSet(cfg, m)
		}
		members = append(members, tm)
	}

	// Pre-compute discriminator entries for ValueByDiscriminator
	var entries []unionTemplateDiscEntry
	for _, m := range cfg.Members {
		for _, dv := range m.DiscriminatorValues {
			entries = append(entries, unionTemplateDiscEntry{Value: dv, MethodName: m.MethodName})
		}
	}

	return unionTemplateData{
		TypeName:             cfg.TypeName,
		Doc:                  cfg.Doc,
		HelperPrefix:         cfg.HelperPrefix,
		FixedFields:          fixedFields,
		Members:              members,
		Discriminator:        cfg.Discriminator,
		DiscriminatorEntries: entries,
	}
}

// allMembersMapped returns true if every union member has at least one discriminator mapping value.
func allMembersMapped(cfg UnionTypeConfig) bool {
	if cfg.Discriminator == nil {
		return false
	}
	for _, m := range cfg.Members {
		if len(m.DiscriminatorValues) == 0 {
			return false
		}
	}
	return true
}

// computeDiscriminatorAutoSet returns the Go statement to auto-set the discriminator
// value in From/Merge methods (e.g., `v.AuthType = "none"` or `t.Type = "v1"`).
func computeDiscriminatorAutoSet(cfg UnionTypeConfig, m UnionMember) string {
	discValue := m.DiscriminatorValues[0]
	propName := cfg.Discriminator.PropertyName

	var goFieldName string
	if cfg.Converter != nil {
		goFieldName = cfg.Converter.ToPropertyName(propName)
	} else {
		goFieldName = UppercaseFirstCharacter(propName)
	}

	if cfg.hasFixedField(propName) {
		return fmt.Sprintf("t.%s = %q", goFieldName, discValue)
	}
	return fmt.Sprintf("v.%s = %q", goFieldName, discValue)
}

// loadUnionTemplate loads and parses the union template.
func loadUnionTemplate() (*template.Template, error) {
	content, err := templates.TemplateFS.ReadFile("files/union/union.go.tmpl")
	if err != nil {
		return nil, fmt.Errorf("reading union template: %w", err)
	}
	funcMap := template.FuncMap{
		"splitLines": func(s string) []string { return strings.Split(s, "\n") },
	}
	tmpl, err := template.New("union").Funcs(funcMap).Parse(string(content))
	if err != nil {
		return nil, fmt.Errorf("parsing union template: %w", err)
	}
	return tmpl, nil
}

// needsTypeConversion returns true if a numeric type needs an explicit conversion
// from the default Go literal type (int for integers, float64 for floats).
func needsTypeConversion(goType string) bool {
	switch goType {
	case "int8", "int16", "int32", "int64",
		"uint", "uint8", "uint16", "uint32", "uint64",
		"float32":
		return true
	default:
		return false
	}
}