// Copyright 2020 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.

package godoc

import (
	"bytes"
	"context"
	"encoding/gob"
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io"
	"path/filepath"
	"reflect"
	"runtime"
	"sort"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
)

var packageToTest string = filepath.Join(runtime.GOROOT(), "src", "net", "http")

func TestEncodeDecodePackage(t *testing.T) {
	p, err := packageForDir(packageToTest, true)
	if err != nil {
		t.Fatal(err)
	}
	var want, got bytes.Buffer
	printPackage(&want, p)
	data, err := p.Encode(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	p2, err := DecodePackage(data)
	if err != nil {
		t.Fatal(err)
	}
	printPackage(&got, p2)
	// Diff the textual output of printPackage, because cmp.Diff takes too long
	// on the Packages themselves.
	if diff := cmp.Diff(want.String(), got.String()); diff != "" {
		t.Errorf("package differs after decoding (-want, +got):\n%s", diff)
	}
}

func TestObjectIdentity(t *testing.T) {
	// Check that encoding and decoding preserves object identity.
	ctx := context.Background()
	const file = `
package p
var a int
func main() { a = 1 }
`

	compareObjs := func(f *ast.File) {
		t.Helper()
		// We know (from hand-inspecting the output of ast.Fprintf) that these two
		// objects are identical in the above program.
		o1 := f.Decls[0].(*ast.GenDecl).Specs[0].(*ast.ValueSpec).Names[0].Obj
		o2 := f.Decls[1].(*ast.FuncDecl).Body.List[0].(*ast.AssignStmt).Lhs[0].(*ast.Ident).Obj
		if o1 != o2 {
			t.Fatal("objects not identical")
		}
	}

	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "test.go", file, parser.ParseComments)
	if err != nil {
		t.Fatal(err)
	}
	compareObjs(f)

	p := NewPackage(fset, "linux", "amd64", nil)
	p.AddFile(f, false)
	data, err := p.Encode(ctx)
	if err != nil {
		t.Fatal(err)
	}
	p, err = DecodePackage(data)
	if err != nil {
		t.Fatal(err)
	}
	compareObjs(p.Files[0].AST)
}

func packageForDir(dir string, removeNodes bool) (*Package, error) {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, dir, nil, parser.ParseComments)
	if err != nil {
		return nil, err
	}
	p := NewPackage(fset, "linux", "amd64", nil)
	for _, pkg := range pkgs {
		for _, f := range pkg.Files {
			p.AddFile(f, removeNodes)
		}
	}
	return p, nil
}

// Compare the time to decode AST files with and without
// removing parts of the AST not relevant to documentation.
//
// Run on a cloudtop 9/29/2020:
// - data size is 3.5x smaller
// - decode time is 4.5x faster
func BenchmarkRemovingAST(b *testing.B) {
	for _, removeNodes := range []bool{false, true} {
		b.Run(fmt.Sprintf("removeNodes=%t", removeNodes), func(b *testing.B) {
			p, err := packageForDir(packageToTest, removeNodes)
			if err != nil {
				b.Fatal(err)
			}
			data, err := p.Encode(context.Background())
			if err != nil {
				b.Fatal(err)
			}
			b.Logf("len(data) = %d", len(data))
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				if _, err := DecodePackage(data); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// printPackage outputs a human-readable form of p to w, deterministically. (The
// ast.Fprint function does not print ASTs deterministically: it is subject to
// random-order map iteration.) The output is designed to be diffed.
func printPackage(w io.Writer, p *Package) error {
	if err := printFileSet(w, p.Fset); err != nil {
		return err
	}
	if _, err := fmt.Fprintf(w, "GOOS %q, GOARCH %q\n", p.GOOS, p.GOARCH); err != nil {
		return err
	}
	var mpps []string
	for k := range p.ModulePackagePaths {
		mpps = append(mpps, k)
	}
	sort.Strings(mpps)
	if _, err := fmt.Fprintf(w, "ModulePackagePaths: %v\n", mpps); err != nil {
		return err
	}

	for _, pf := range p.Files {
		if _, err := fmt.Fprintf(w, "---- %s\n", pf.Name); err != nil {
			return err
		}
		if err := printNode(w, pf.AST); err != nil {
			return err
		}
	}
	return nil
}

func printNode(w io.Writer, root ast.Node) error {
	var err error
	seen := map[interface{}]int{}

	pr := func(format string, args ...interface{}) {
		if err == nil {
			_, err = fmt.Fprintf(w, format, args...)
		}
	}

	indent := func(d int) {
		for i := 0; i < d; i++ {
			pr("  ")
		}
	}

	var prValue func(interface{}, int)
	prValue = func(x interface{}, depth int) {
		indent(depth)
		if x == nil || reflect.ValueOf(x).IsNil() {
			pr("nil\n")
			return
		}
		ts := strings.TrimPrefix(fmt.Sprintf("%T", x), "*ast.")
		if idx, ok := seen[x]; ok {
			pr("%s@%d\n", ts, idx)
			return
		}
		idx := len(seen)
		seen[x] = idx
		pr("%s#%d", ts, idx)
		if obj, ok := x.(*ast.Object); ok {
			pr(" %s %s %v\n", obj.Name, obj.Kind, obj.Data)
			prValue(obj.Decl, depth+1)
			return
		}
		n, ok := x.(ast.Node)
		if !ok {
			pr(" %v\n", x)
			return
		}
		pr(" %d-%d", n.Pos(), n.End())
		switch n := n.(type) {
		case *ast.Ident:
			pr(" %q\n", n.Name)
			if n.Obj != nil {
				prValue(n.Obj, depth+1)
			}
		case *ast.BasicLit:
			pr(" %s %s %d\n", n.Value, n.Kind, n.ValuePos)
		case *ast.UnaryExpr:
			pr(" %s\n", n.Op)
		case *ast.BinaryExpr:
			pr(" %s\n", n.Op)
		case *ast.Comment:
			pr(" %q\n", n.Text)
		case *ast.File:
			// Doc, Name and Decls are walked, but not Scope or Unresolved.
			if n.Scope != nil {
				pr(" Scope.Outer: %p\n", n.Scope.Outer)
				var keys []string
				for k := range n.Scope.Objects {
					keys = append(keys, k)
				}
				sort.Strings(keys)
				for _, k := range keys {
					pr("  key %q\n", k)
					prValue(n.Scope.Objects[k], depth+1)
				}
			}
			indent(depth)
			pr("unresolved:\n")
			for _, id := range n.Unresolved {
				prValue(id, depth+1)
			}
		default:
			pr("\n")
		}
		ast.Inspect(n, func(m ast.Node) bool {
			if m == n {
				return true
			}
			if m != nil {
				prValue(m, depth+1)
			}
			return false
		})
	}

	prValue(root, 0)
	return err
}

// Compare the time it takes to encode with gob vs. internal/codec.
func BenchmarkEncoding(b *testing.B) {
	p, err := packageForDir(packageToTest, true)
	if err != nil {
		b.Fatal(err)
	}
	b.Run("gob", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			_, err := p.gobEncode()
			if err != nil {
				b.Fatal(err)
			}
		}
	})
	b.Run("fast", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			_, err := p.fastEncode()
			if err != nil {
				b.Fatal(err)
			}
		}
	})
}

// Compare the time it takes to decode with gob vs. internal/codec.
func BenchmarkDecoding(b *testing.B) {
	p, err := packageForDir(packageToTest, true)
	if err != nil {
		b.Fatal(err)
	}
	b.Run("gob", func(b *testing.B) {
		data, err := p.gobEncode()
		if err != nil {
			b.Fatal(err)
		}
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			_, err := gobDecodePackage(data[encodingTypeLen:])
			if err != nil {
				b.Fatal(err)
			}
		}
	})
	b.Run("fast", func(b *testing.B) {
		data, err := p.fastEncode()
		if err != nil {
			b.Fatal(err)
		}
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			_, err := fastDecodePackage(data[encodingTypeLen:])
			if err != nil {
				b.Fatal(err)
			}
		}
	})
}

func printFileSet(w io.Writer, fset *token.FileSet) error {
	var err error
	fset.Iterate(func(f *token.File) bool {
		_, err = fmt.Fprintf(w, "%s %d %d %d\n", f.Name(), f.Base(), f.Size(), f.LineCount())
		return err == nil
	})
	return err
}

func (p *Package) gobEncode() (_ []byte, err error) {
	if p.renderCalled {
		return nil, errors.New("can't Encode after Render")
	}

	for _, f := range p.Files {
		removeCycles(f)
	}

	var buf bytes.Buffer
	io.WriteString(&buf, gobEncodingType)
	enc := gob.NewEncoder(&buf)
	// Encode the fset using the Write method it provides.
	if err := p.Fset.Write(enc.Encode); err != nil {
		return nil, fmt.Errorf("p.Fset.Write: %v", err)
	}
	if err := enc.Encode(p.encPackage); err != nil {
		return nil, fmt.Errorf("enc.Encode: %v", err)
	}
	for _, f := range p.Files {
		fixupObjects(f)
	}
	return buf.Bytes(), nil
}