package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"sort"
	"strings"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s <output-file>\n", os.Args[0])
		os.Exit(1)
	}

	outputFile := os.Args[1]

	// Parse the spec.go file to extract Source field information
	sourceFields, err := extractSourceFields()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error extracting source fields: %v\n", err)
		os.Exit(1)
	}

	// Generate the code
	code, err := generateCode(sourceFields)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error generating code: %v\n", err)
		os.Exit(1)
	}

	// Write to output file
	err = os.WriteFile(outputFile, code, 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing output file: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Generated source variant methods in %s\n", outputFile)
}

type SourceField struct {
	Name     string // Field name (e.g. "DockerImage")
	TypeName string // Type name (e.g. "SourceDockerImage")
}

func extractSourceFields() ([]SourceField, error) {
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, "spec.go", nil, parser.ParseComments)
	if err != nil {
		return nil, fmt.Errorf("failed to parse spec.go: %w", err)
	}

	var sourceFields []SourceField

	// Find the Source struct
	ast.Inspect(node, func(n ast.Node) bool {
		if ts, ok := n.(*ast.TypeSpec); ok && ts.Name.Name == "Source" {
			if st, ok := ts.Type.(*ast.StructType); ok {
				for _, field := range st.Fields.List {
					if len(field.Names) == 0 {
						continue // Skip embedded fields
					}

					fieldName := field.Names[0].Name

					// Check if this is a pointer to a type starting with "Source"
					if ptr, ok := field.Type.(*ast.StarExpr); ok {
						if ident, ok := ptr.X.(*ast.Ident); ok {
							typeName := ident.Name
							if strings.HasPrefix(typeName, "Source") {
								sourceFields = append(sourceFields, SourceField{
									Name:     fieldName,
									TypeName: typeName,
								})
							}
						}
					}
				}
			}
		}
		return true
	})

	// Sort fields for consistent output
	sort.Slice(sourceFields, func(i, j int) bool {
		return sourceFields[i].Name < sourceFields[j].Name
	})

	return sourceFields, nil
}

func generateCode(fields []SourceField) ([]byte, error) {
	// Sort fields to ensure consistent output
	sort.Slice(fields, func(i, j int) bool {
		return fields[i].Name < fields[j].Name
	})

	var buf bytes.Buffer

	// Generate file header
	buf.WriteString(`// Code generated by cmd/gen-source-variants. DO NOT EDIT.

package dalec

import (
	"fmt"
)

`)

	// Generate validateSourceVariants method
	buf.WriteString("// validateSourceVariants ensures exactly one source variant is set\n")
	buf.WriteString("func (s *Source) validateSourceVariants() error {\n")
	buf.WriteString("\tcount := 0\n\n")

	for _, field := range fields {
		buf.WriteString(fmt.Sprintf("\tif s.%s != nil {\n", field.Name))
		buf.WriteString("\t\tcount++\n")
		buf.WriteString("\t}\n")
	}

	buf.WriteString(`
	switch count {
	case 0:
		return fmt.Errorf("no non-nil source variant")
	case 1:
		return nil
	default:
		return fmt.Errorf("more than one source variant defined")
	}
}

`)

	// Generate toInterface method
	buf.WriteString("// toInterface returns the underlying source interface implementation\n")
	buf.WriteString("func (s *Source) toInterface() source {\n")
	buf.WriteString("\tswitch {\n")

	for _, field := range fields {
		buf.WriteString(fmt.Sprintf("\tcase s.%s != nil:\n", field.Name))
		buf.WriteString(fmt.Sprintf("\t\treturn s.%s\n", field.Name))
	}

	buf.WriteString(`	default:
		panic(errNoSourceVariant)
	}
}
`)

	// Format the generated code
	formatted, err := format.Source(buf.Bytes())
	if err != nil {
		return nil, fmt.Errorf("failed to format generated code: %w", err)
	}

	return formatted, nil
}
