// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The bundle command concatenates the source files of a package,
// renaming package-level names by adding a prefix and renaming
// identifiers as needed to preserve referential integrity.
//
// Example:
// 	$ bundle golang.org/x/net/http2 net/http http2
//
// The command above prints a single file containing the code of
// golang.org/x/net/http2, suitable for inclusion in package net/http,
// in which toplevel names have been prefixed with "http2".
//
// Assumptions:
// - no file in the package imports "C", that is, uses cgo.
// - no file in the package has GOOS or GOARCH build tags or file names.
// - comments associated with the package or import declarations,
//   may be discarded, as may comments associated with no top-level
//   declaration at all.
// - neither the original package nor the destination package contains
//   any identifiers starting with the designated prefix.
//   (This allows us to avoid various conflict checks.)
// - there are no renaming imports.
// - test files are ignored.
// - none of the renamed identifiers is significant
//   to reflection-based logic.
//
// Only package-level var, func, const, and type objects are renamed,
// and embedded fields of renamed types.  No methods are renamed, so we
// needn't worry about preserving interface satisfaction.
//
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
)

func main() {
	log.SetPrefix("bundle: ")
	log.SetFlags(0)

	flag.Parse()
	args := flag.Args()
	if len(args) != 3 {
		log.Fatal(`Usage: bundle package dest prefix

Arguments:
 package is the import path of the package to concatenate.
 dest is the import path of the package in which the resulting file will reside.
 prefix is the string to attach to all renamed identifiers.
`)
	}
	initialPkg, dest, prefix := args[0], args[1], args[2]

	if err := bundle(os.Stdout, initialPkg, dest, prefix); err != nil {
		log.Fatal(err)
	}
}

var ctxt = &build.Default

func bundle(w io.Writer, initialPkg, dest, prefix string) error {
	// Load the initial package.
	conf := loader.Config{ParserMode: parser.ParseComments, Build: ctxt}
	conf.TypeCheckFuncBodies = func(p string) bool { return p == initialPkg }
	conf.Import(initialPkg)

	lprog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}

	info := lprog.Package(initialPkg)

	objsToUpdate := make(map[types.Object]bool)
	var rename func(from types.Object)
	rename = func(from types.Object) {
		if !objsToUpdate[from] {
			objsToUpdate[from] = true

			// Renaming a type that is used as an embedded field
			// requires renaming the field too. e.g.
			// 	type T int // if we rename this to U..
			// 	var s struct {T}
			// 	print(s.T) // ...this must change too
			if _, ok := from.(*types.TypeName); ok {
				for id, obj := range info.Uses {
					if obj == from {
						if field := info.Defs[id]; field != nil {
							rename(field)
						}
					}
				}
			}
		}
	}

	// Rename each package-level object.
	scope := info.Pkg.Scope()
	for _, name := range scope.Names() {
		rename(scope.Lookup(name))
	}

	var out bytes.Buffer

	fmt.Fprintf(&out, "// Code generated by golang.org/x/tools/cmd/bundle command:\n")
	fmt.Fprintf(&out, "//   $ bundle %s %s %s\n\n", initialPkg, dest, prefix)

	// Concatenate package comments from all files...
	for _, f := range info.Files {
		if doc := f.Doc.Text(); strings.TrimSpace(doc) != "" {
			for _, line := range strings.Split(doc, "\n") {
				fmt.Fprintf(&out, "// %s\n", line)
			}
		}
	}
	// ...but don't let them become the actual package comment.
	fmt.Fprintln(&out)

	// TODO(adonovan): don't assume pkg.name == basename(pkg.path).
	fmt.Fprintf(&out, "package %s\n\n", filepath.Base(dest))

	// Print a single declaration that imports all necessary packages.
	// TODO(adonovan):
	// - support renaming imports.
	// - preserve comments from the original import declarations.
	for _, f := range info.Files {
		for _, imp := range f.Imports {
			if imp.Name != nil {
				log.Fatalf("%s: renaming imports not supported",
					lprog.Fset.Position(imp.Pos()))
			}
		}
	}
	fmt.Fprintln(&out, "import (")
	for _, p := range info.Pkg.Imports() {
		if p.Path() == dest {
			continue
		}
		fmt.Fprintf(&out, "\t%q\n", p.Path())
	}
	fmt.Fprintln(&out, ")\n")

	// Modify and print each file.
	for _, f := range info.Files {
		// Update renamed identifiers.
		for id, obj := range info.Defs {
			if objsToUpdate[obj] {
				id.Name = prefix + obj.Name()
			}
		}
		for id, obj := range info.Uses {
			if objsToUpdate[obj] {
				id.Name = prefix + obj.Name()
			}
		}

		// For each qualified identifier that refers to the
		// destination package, remove the qualifier.
		// The "@@@." strings are removed in postprocessing.
		ast.Inspect(f, func(n ast.Node) bool {
			if sel, ok := n.(*ast.SelectorExpr); ok {
				if id, ok := sel.X.(*ast.Ident); ok {
					if obj, ok := info.Uses[id].(*types.PkgName); ok {
						if obj.Imported().Path() == dest {
							id.Name = "@@@"
						}
					}
				}
			}
			return true
		})

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		//
		// TODO(adonovan): this may cause loss of comments
		// preceding or associated with the package or import
		// declarations or not associated with any declaration.
		// Check.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			buf.Reset()
			format.Node(&buf, lprog.Fset, decl)
			// Remove each "@@@." in the output.
			// TODO(adonovan): not hygienic.
			out.Write(bytes.Replace(buf.Bytes(), []byte("@@@."), nil, -1))
			out.WriteString("\n\n")
		}
	}

	// Now format the entire thing.
	result, err := format.Source(out.Bytes())
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}

	_, err = w.Write(result)
	return err
}