// Copyright 2025 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
 "bytes"
 "errors"
 "flag"
 "fmt"
 "go/ast"
 "go/format"
 "go/parser"
 "go/scanner"
 "go/token"
 "maps"
 "os"
 "path"
 "path/filepath"
 "slices"
 "sort"
 "strings"
)

// This file provides functionality to auto-generate custom Encode and Decode
// method for Go structs.
//
// # Auto-Generation Trigger
//
// The generation process is triggered for structs annotated with the comment:
//   // @auto-generate: gob
// within a given Go source file.
//
// # Generated Methods
//
// 1. Encode(buf *bytes.Buffer) error:
//    Encodes each field within the struct into a stream of bytes.
//
// 2. Decode(buf *bytes.Reader) error:
//    Decodes from a stream of byte and populates each field within the struct.
//
// # Supported Data Types
//
// The generator can handle the following data types within the annotated structs:
//   - Most of the primitive types (e.g., int, string, bool)
//   - Pointers to supported types
//   - Maps (keys and values must be supported types)
//   - Slices of supported types
//   - Type aliases
//   - Interfaces.
//   - Other user defined structs (which should have generated Encode and Decode methods)

var verify = flag.Bool("verify", false, "verify existing outputs")

const genGobAnnotation = "@auto-generate: gob"
const blueprintPkgPrefix = "github.com/google/blueprint"
const blueprintPkgPath = "build/blueprint"
const soongPkgPrefix = "android/soong"
const soongPkgPath = "build/soong"
const gobtoolsImport = `"github.com/google/blueprint/gobtools"`
const proptoolsImport = `"github.com/google/blueprint/proptools"`
const reflectImport = `"reflect"`
const valueIsNil = -1
const codeGenComment = "// Code generated by go run gob_gen.go; DO NOT EDIT."

type typeDefTypes int

const (
 Struct typeDefTypes = iota
 Alias
 Interface
 Slice
 Array
 Map
 Pointer
 Ident
 Unknown
)

type structMap map[string]*ast.TypeSpec

type gobGen struct {
 pkgStructs map[string]structMap
 importPkgs map[string]string
 curPackage string
 sourceDir  string
 fieldId    int
 imports    map[string]bool
}

func newGobGen() *gobGen {
 return &gobGen{
  pkgStructs: make(map[string]structMap),
  importPkgs: make(map[string]string),
  imports:    make(map[string]bool),
 }
}

func main() {
 flag.Usage = func() {
  fmt.Fprintf(os.Stderr, "Usage: %s [sources]\n", os.Args[0])
  flag.PrintDefaults()
 }
 flag.Parse()
 sources := slices.Clone(flag.Args())

 if f := os.Getenv("GOFILE"); f != "" {
  sources = append(sources, f)
 }

 if len(sources) == 0 {
  flag.Usage()
 }

 g := newGobGen()
 curDir, _ := os.Getwd()
 if *verify {
  // verify mode, the current directory is the base of the source tree.
  g.sourceDir = curDir
 } else {
  var err error
  if g.sourceDir, err = getTop(curDir); err != nil {
   panic(err)
  }
 }

 var out []byte
 var outputFile string
 var err error
 for _, s := range sources {
  out, outputFile, err = g.generate(s, out, *verify)
  if err != nil {
   fmt.Fprintf(os.Stderr, "failed to generate output for %s: %s\n", s, err)
   os.Exit(1)
  }
 }

 if *verify {
  if len(out) == 0 {
   err := expectNotExist(outputFile)
   if err != nil {
    fmt.Fprintf(os.Stderr, "verification error: %s\n", err)
    os.Exit(1)
   }
  } else {
   if err := expectContents(outputFile, out); err != nil {
    fmt.Fprintf(os.Stderr, "verification error: %s\n", err)
    os.Exit(1)
   }
   if !slices.Contains(sources, outputFile) {
    fmt.Fprintf(os.Stderr, "verification error: generated file %s is not in srcs\n", outputFile)
    os.Exit(1)
   }
  }
 } else if len(out) > 0 {
  err = os.WriteFile(outputFile, out, 0666)
  if err != nil {
   fmt.Fprintf(os.Stderr, "failed to write output to %s: %s\n", outputFile, err)
   os.Exit(1)
  }
 }
}

const (
 beginMarkPrefix = "// begin of "
 endMarkPrefix   = "// end of "
)

// mergeSections merges the newly generated section with the old ones in alphabetical order.
func mergeSections(origContent []byte, newContent string) []byte {
 // Combine old and new sections. This ensures all sections and imports are processed together.
 combinedContent := string(origContent) + "\n" + newContent

 // Perform a pass over the combined content to parse everything
 var packageLine, genComment string
 allImportsSet := make(map[string]struct{})
 sectionsMap := make(map[string]string)

 inImportBlock := false
 inSectionBlock := false
 var currentSectionMarker string
 var currentSectionContent []string

 for _, line := range strings.Split(combinedContent, "\n") {
  trimmedLine := strings.TrimSpace(line)

  // Capture the first valid header line found
  if packageLine == "" && strings.HasPrefix(trimmedLine, "package ") {
   packageLine = trimmedLine
   continue
  }
  if genComment == "" && strings.HasPrefix(trimmedLine, codeGenComment) {
   genComment = trimmedLine
   continue
  }

  // Collect all unique imports
  if strings.HasPrefix(trimmedLine, "import (") {
   inImportBlock = true
   continue
  }
  if inImportBlock {
   if trimmedLine == ")" {
    inImportBlock = false
   } else if trimmedLine != "" {
    allImportsSet[trimmedLine] = struct{}{}
   }
   continue
  }
  if strings.HasPrefix(trimmedLine, "import ") {
   path := strings.Trim(trimmedLine[len("import "):], `"`)
   allImportsSet[`"`+path+`"`] = struct{}{}
   continue
  }

  // Parse generated sections. The map naturally handles updates by overwriting keys.
  if strings.HasPrefix(trimmedLine, beginMarkPrefix) {
   inSectionBlock = true
   currentSectionMarker = trimmedLine
   currentSectionContent = nil // Reset for new section
   continue
  }
  if inSectionBlock {
   if strings.HasPrefix(trimmedLine, endMarkPrefix) {
    inSectionBlock = false
    sectionsMap[currentSectionMarker] = strings.Join(currentSectionContent, "\n")
   } else {
    // Don't append blank lines at the start of the content
    if len(currentSectionContent) > 0 || trimmedLine != "" {
     currentSectionContent = append(currentSectionContent, line)
    }
   }
  }
 }

 // Rebuild the entire file from the collected parts.
 var output strings.Builder

 // Write header
 if genComment != "" {
  output.WriteString(genComment + "\n\n")
 }
 if packageLine != "" {
  output.WriteString(packageLine + "\n\n")
 }

 // Write sorted imports
 if len(allImportsSet) > 0 {
  var allImports []string
  for imp := range allImportsSet {
   allImports = append(allImports, imp)
  }
  sort.Strings(allImports)

  output.WriteString("import (\n")
  for _, imp := range allImports {
   output.WriteString("\t" + imp + "\n")
  }
  output.WriteString(")\n")
 }

 // Get all section markers and sort them alphabetically
 var sortedMarkers []string
 for marker := range sectionsMap {
  sortedMarkers = append(sortedMarkers, marker)
 }
 sort.Strings(sortedMarkers)

 // Write all sections back in sorted order
 for _, marker := range sortedMarkers {
  content := strings.TrimSpace(sectionsMap[marker])
  endMarker := strings.Replace(marker, beginMarkPrefix, endMarkPrefix, 1)

  output.WriteString("\n" + marker + "\n")
  if content != "" {
   output.WriteString(content + "\n\n")
  }
  output.WriteString(endMarker + "\n")
 }

 finalString := strings.TrimRight(output.String(), "\n") + "\n"
 return []byte(finalString)
}

// getTop finds the root of the source tree.
// It mimics the behavior of the build/make/envsetup.sh.
func getTop(curDir string) (string, error) {
 const topFile = "build/make/core/envsetup.mk"

 // Start searching from the current directory.
 dir := curDir
 for {
  // Check for the existence of the marker file in the current directory.
  pathToCheck := filepath.Join(dir, topFile)
  if info, err := os.Stat(pathToCheck); err == nil && !info.IsDir() {
   // Found it. Resolve symlinks and return the path.
   return filepath.EvalSymlinks(dir)
  }

  // Get the parent directory.
  parentDir := filepath.Dir(dir)

  // If the parent is the same as the current directory, we've hit the root.
  if parentDir == dir {
   break
  }

  // Move up to the parent directory for the next iteration.
  dir = parentDir
 }

 // If the loop completes without finding the file, report error.
 panic(fmt.Errorf("failed to find the root of the source tree: %s", curDir))
}

func (g *gobGen) findPackagePath(pkgName string) string {
 var pkgDir string
 baseDir := strings.TrimPrefix(g.importPkgs[pkgName], blueprintPkgPrefix)
 if baseDir != g.importPkgs[pkgName] {
  pkgDir = filepath.Join(g.sourceDir, blueprintPkgPath, baseDir)
 } else {
  baseDir = strings.TrimPrefix(g.importPkgs[pkgName], soongPkgPrefix)
  if baseDir != g.importPkgs[pkgName] {
   pkgDir = filepath.Join(g.sourceDir, soongPkgPath, baseDir)
  }
 }
 if pkgDir == "" {
  panic(fmt.Errorf("failed to find the package path: %s %s", pkgName, g.importPkgs[pkgName]))
 }

 return pkgDir
}

func (g *gobGen) importPackage(pkgName string, pkgDir string) {
 if _, ok := g.pkgStructs[pkgName]; ok {
  return
 }
 includePattern := "*.go"
 excludePattern := "*_test.go"

 fullPattern := filepath.Join(pkgDir, includePattern)

 matches, err := filepath.Glob(fullPattern)
 if err != nil {
  panic(fmt.Errorf("error matching include pattern: %v", err))
 }

 if len(matches) == 0 {
  panic(fmt.Errorf("No files found matching pattern '%s' in directory '%s'\n", includePattern, pkgDir))
  return
 }

 for _, match := range matches {
  ok, err := filepath.Match(excludePattern, match)
  if err != nil {
   panic(fmt.Errorf("error matching exclude pattern: %v", err))
  }
  if ok {
   continue
  }
  node, err := parseFile(match)
  if err != nil {
   panic(fmt.Errorf("failed to parse file: %s", match))
  }
  g.loadTypes(node)
 }
}

func parseFile(source string) (*ast.File, error) {
 fset := token.NewFileSet()
 node, err := parser.ParseFile(fset, source, nil, parser.ParseComments)
 if err != nil {
  // ParseFile might return multiple errors in the form of a scanner.ErrorList.  By default printing the error
  // only shows the first error.  Verification may happen very early during the build, so this may be the first
  // time syntax errors are reported.  Use scanner.PrintError to convert them into a single error that contains
  // all the error lines to make the errors more actionable.
  if errorList, ok := err.(scanner.ErrorList); ok {
   var buf bytes.Buffer
   scanner.PrintError(&buf, errorList)
   err = errors.New(buf.String())
  }
  return nil, fmt.Errorf("failed to parse:\n%w", err)
 }
 return node, nil
}

func (g *gobGen) generateRegistry(structDecl *ast.TypeSpec, codeBody *strings.Builder, initCodeBody *strings.Builder) {
 structName := structDecl.Name.Name
 typeId := structName + "GobRegId"
 codeBody.WriteString(fmt.Sprintf("\tvar %s int16\n", typeId))
 codeBody.WriteString("func (r " + structName + ") GetTypeId() int16 {\n")
 codeBody.WriteString(fmt.Sprintf("\treturn %s\n", typeId))
 codeBody.WriteString("}\n\n")

 initCodeBody.WriteString(fmt.Sprintf("\t%s = gobtools.RegisterType(func() gobtools.CustomDec { return new(%s) })\n", typeId, structName))
 g.imports[gobtoolsImport] = true
}

func (g *gobGen) generate(source string, input []byte, verify bool) ([]byte, string, error) {
 node, err := parseFile(source) // Find the file containing the struct
 if err != nil {
  return nil, "", err
 }
 for _, p := range node.Imports {
  pkgPath := strings.ReplaceAll(p.Path.Value, "\"", "")
  if strings.HasPrefix(pkgPath, blueprintPkgPrefix) || strings.HasPrefix(pkgPath, soongPkgPrefix) {
   pkgName := path.Base(pkgPath)
   if _, ok := g.importPkgs[pkgName]; !ok {
    g.importPkgs[pkgName] = pkgPath
   }
  }
 }
 g.curPackage = node.Name.Name
 g.importPackage(g.curPackage, path.Dir(source))

 var b bytes.Buffer
 fmt.Fprintf(&b, "%s\n\n", codeGenComment)
 fmt.Fprintf(&b, "package %s\n", g.curPackage)
 fmt.Fprintln(&b, "import (")

 var codeBodies []*strings.Builder
 g.imports = make(map[string]bool)
 initCodeBody := &strings.Builder{}
 initCodeBody.WriteString("func init() {\n")
 structDecls := findStructs(node)
 for _, structDecl := range structDecls {
  g.fieldId = 0
  codeBody := &strings.Builder{}
  g.generateEncode(g.curPackage, structDecl, codeBody)
  codeBody.WriteString("\n")
  g.fieldId = 0
  g.generateHash(g.curPackage, structDecl, codeBody)
  codeBody.WriteString("\n")
  g.fieldId = 0
  g.generateDecode(g.curPackage, structDecl, codeBody)
  codeBodies = append(codeBodies, codeBody)
  g.generateRegistry(structDecl, codeBody, initCodeBody)
 }

 outputFile := filepath.Join(filepath.Dir(source), g.curPackage+"_enc.go")
 if len(codeBodies) == 0 {
  return input, outputFile, nil
 }

 initCodeBody.WriteString("}\n\n")

 fmt.Fprintln(&b, strings.Join(slices.Sorted(maps.Keys(g.imports)), "\n"))
 fmt.Fprintln(&b, ")")
 beginMark := beginMarkPrefix + filepath.Base(source)
 endMark := endMarkPrefix + filepath.Base(source)
 fmt.Fprintln(&b, beginMark)
 fmt.Fprintf(&b, initCodeBody.String())
 for _, codeBody := range codeBodies {
  fmt.Fprintf(&b, codeBody.String())
  fmt.Fprintln(&b)
 }
 fmt.Fprintln(&b, endMark)

 out, err := format.Source(b.Bytes())
 if err != nil {
  return nil, "", fmt.Errorf("source format error: %w", err)
 }

 if !verify {
  input, err = os.ReadFile(outputFile)
  if err != nil && !os.IsNotExist(err) {
   return nil, "", err
  }
 }
 out = mergeSections(input, string(out))
 return out, outputFile, nil
}

// expectContents verifies the that file contains the given bytes, returning an error that describes
// how to fix the problem if it does not.
func expectContents(file string, expected []byte) error {
 actual, err := os.ReadFile(file)
 if os.IsNotExist(err) {
  return fmt.Errorf("generated file %s does not exist, rerun `go generate` in %s",
   file, filepath.Dir(file))
 }
 if err != nil {
  return err
 }

 if len(expected) == 0 {
  return fmt.Errorf("found unexpected generated file %s, delete it", file)
 }

 if !bytes.Equal(actual, expected) {
  missing := missingSections(actual, expected)
  if len(missing) > 0 {
   return fmt.Errorf("some files were removed: %s, delete %s and rerun `go generate` in %s",
    strings.Join(missing, ", "), file, filepath.Dir(file))
  }
  return fmt.Errorf("generated file %s has out of date contents, rerun `go generate` in %s",
   file, filepath.Dir(file))
 }

 return nil
}

// parseSectionMarkers scans content and returns a set of all begin markers found.
func parseSectionMarkers(content []byte) map[string]struct{} {
 markers := make(map[string]struct{})
 lines := strings.Split(string(content), "\n")

 for _, line := range lines {
  trimmedLine := strings.TrimSpace(line)
  if strings.HasPrefix(trimmedLine, beginMarkPrefix) {
   markers[trimmedLine] = struct{}{}
  }
 }
 return markers
}

// missingSections finds the missing sections in new content from the original one.
// It returns missing sections.
func missingSections(origContent, newContent []byte) (missing []string) {
 // Step 1: Parse all markers from each input into sets.
 origSections := parseSectionMarkers(origContent)
 newSections := parseSectionMarkers(newContent)

 for marker := range origSections {
  if _, found := newSections[marker]; !found {
   missing = append(missing, strings.TrimPrefix(marker, beginMarkPrefix))
  }
 }

 return missing
}

// expectNotExist verifies that the file does not exist, returning an error that describes how to
// fix the problem if it does.
func expectNotExist(file string) error {
 _, err := os.Stat(file)
 if os.IsNotExist(err) {
  return nil
 }
 if err != nil {
  return err
 }
 return fmt.Errorf("expected %s to not exist, delete it", file)
}