// Copyright (C) 2016 The Android Open Source Project
//
// 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 (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"time"

	"android.googlesource.com/platform/tools/gpu/client/adb"
	"android.googlesource.com/platform/tools/gpu/client/android"
	"android.googlesource.com/platform/tools/gpu/client/android/apk"
	"android.googlesource.com/platform/tools/gpu/client/bind"
	"android.googlesource.com/platform/tools/gpu/client/gapii"
	"android.googlesource.com/platform/tools/gpu/client/process"
	"android.googlesource.com/platform/tools/gpu/framework/app"
	"android.googlesource.com/platform/tools/gpu/framework/device"
	"android.googlesource.com/platform/tools/gpu/framework/file"
	"android.googlesource.com/platform/tools/gpu/framework/log"
	"android.googlesource.com/platform/tools/gpu/framework/task"
	"android.googlesource.com/platform/tools/gpu/gapid/pkgdata"
)

var (
	targetDevice    string
	duration        time.Duration
	output          string
	localPort       int
	localApp        file.Path
	observeFrames   uint
	observeDraws    uint
	disablePCS      bool
	recordErrors    bool
	clearCache      bool
	inputFile       string
	recordInputs    bool
	replayInputs    bool
	useLDPreload    bool
	apkPath         file.Path
	interceptorPath file.Path
	gapiiPath       file.Path
	vulkanLayerPath file.Path
)

func init() {
	verb := &app.Verb{
		Name:      "trace",
		ShortHelp: "Captures a gfx trace from an application",
		Run:       doTrace,
	}
	verb.Flags.StringVar(&targetDevice, "device", "", "the device to capture on")
	verb.Flags.DurationVar(&duration, "d", 0, "duration to trace for")
	verb.Flags.StringVar(&output, "out", "", "the file to generate")
	verb.Flags.IntVar(&localPort, "local", 0, "capture a local program instead of using ADB")
	verb.Flags.Var(&localApp, "local-app", "a local program to trace")
	verb.Flags.UintVar(&observeFrames, "observe-frames", 0, "capture the framebuffer every n frames (0 to disable)")
	verb.Flags.UintVar(&observeDraws, "observe-draws", 0, "capture the framebuffer every n draws (0 to disable)")
	verb.Flags.BoolVar(&disablePCS, "disable-pcs", false, "disable pre-compiled shaders")
	verb.Flags.BoolVar(&recordErrors, "record-errors", false, "record device error state")
	verb.Flags.BoolVar(&clearCache, "clear-cache", false, "clear package data before running it")
	verb.Flags.StringVar(&inputFile, "input-file", "", "the file to use for recorded inputs")
	verb.Flags.BoolVar(&recordInputs, "record-inputs", false, "record the inputs to file")
	verb.Flags.BoolVar(&replayInputs, "replay-inputs", false, "replay the inputs from file")
	verb.Flags.BoolVar(&useLDPreload, "ld-preload", false, "use LD_PRELOAD to register the interceptor library")
	verb.Flags.Var(&apkPath, "apk", "the path to an apk to install")
	verb.Flags.Var(&interceptorPath, "interceptor", "the path to the interceptor library")
	verb.Flags.Var(&gapiiPath, "gapii", "the path to the gapii library")
	verb.Flags.Var(&vulkanLayerPath, "vulkan-layer", "the path to the vulkan trace layer library")
	app.AddVerb(verb)
}

func doTrace(ctx log.Context, flags flag.FlagSet) error {
	options := gapii.Options{
		ObserveFrameFrequency: uint32(observeFrames),
		ObserveDrawFrequency:  uint32(observeDraws),
		APK:                   apkPath,
		Interceptor:           interceptorPath,
		Gapii:                 gapiiPath,
		VkLayer:               vulkanLayerPath,
	}

	if disablePCS {
		options.Flags |= gapii.DisablePrecompiledShaders
	}
	if recordErrors {
		options.Flags |= gapii.RecordErrorState
	}

	if !localApp.IsEmpty() {
		// Run the local application with VK_LAYER_PATH, VK_INSTANCE_LAYERS,
		// VK_DEVICE_LAYERS and LD_PRELOAD set to correct values to load the spy
		// layer.
		localAbi := device.UnknownABI
		switch runtime.GOOS {
		case "linux":
			localAbi = device.LinuxX86_64
		default:
			return fmt.Errorf("Unsupported OS for local tracing")
		}

		var err error
		if vulkanLayerPath.IsEmpty() {
			vulkanLayerPath, err = pkgdata.File(ctx, localAbi, "LibVkLayerGraphicsSpy.so")
		}
		if err != nil {
			return err
		}
		if gapiiPath.IsEmpty() {
			// TODO (qining): library name may change for different OS/ABI
			gapiiPath, err = pkgdata.File(ctx, localAbi, "libgapii.so")
		}
		if err != nil {
			return err
		}
		libGapiiDirStr := filepath.Dir(gapiiPath.String())
		libGapiiPathStr := filepath.Join(libGapiiDirStr, "libgapii.so")
		vulkanLayerDirStr := filepath.Dir(vulkanLayerPath.String())

		os.Setenv("VK_LAYER_PATH", vulkanLayerDirStr)
		os.Setenv("VK_INSTANCE_LAYERS", "VkGraphicsSpy")
		os.Setenv("VK_DEVICE_LAYERS", "VkGraphicsSpy")
		os.Setenv("LD_PRELOAD", libGapiiPathStr)
		boundPort, err := process.Start(ctx, localApp.String())
		if err != nil {
			return err
		}
		if localPort == 0 {
			localPort = boundPort
		}
	}

	if localPort != 0 {
		return captureLocal(ctx, flags, localPort, options)
	}

	return captureADB(ctx, flags, options, useLDPreload)
}

func captureLocal(ctx log.Context, flags flag.FlagSet, port int, options gapii.Options) error {
	if output == "" {
		output = "capture.gfxtrace"
	}
	return capture(ctx, options, port, output)
}

// gapiiABI returns the ABI-flavor of GAPII to use given the device and package.
func gapiiABI(ctx log.Context, d bind.Device, info *apk.Information) device.ABI {
	for _, preferred := range d.Info().Configuration.ABIs {
		for _, abi := range info.ABI {
			if device.SameABI(preferred, *abi) {
				return preferred
			}
		}
	}
	return d.ABI() // Meh.
}

func captureADB(ctx log.Context, flags flag.FlagSet, options gapii.Options, useLDPreload bool) error {
	d, err := getADBDevice(ctx, targetDevice)
	if err != nil {
		return err
	}
	var a *android.Action
	if !options.APK.IsEmpty() {
		ctx.V("Source", options.APK).Print("Installing APK")
		data, err := ioutil.ReadFile(options.APK.System())
		if err != nil {
			return ctx.WrapError(err, "Read APK")
		}
		info, err := apk.Analyze(ctx, data)
		if err != nil {
			return ctx.WrapError(err, "Analyse APK")
		}
		if err := d.InstallAPK(ctx, options.APK.System(), true, true); err != nil {
			return ctx.WrapError(err, "Install APK")
		}
		pkg := &android.InstalledPackage{
			Name:       info.Name,
			Device:     d,
			ABI:        gapiiABI(ctx, d, info),
			Debuggable: info.Debuggable,
		}
		defer func() {
			ctx.Notice().V("apk", options.APK).Log("Uninstall APK")
			pkg.Uninstall(ctx)
		}()
		a = &android.Action{
			Name:     info.Action,
			Package:  pkg,
			Activity: info.Activity,
		}
	} else {
		if flags.NArg() != 1 {
			app.Usage(ctx, "Invalid number of arguments. Expected 1, got %d", flags.NArg())
			return nil
		}
		activity := flags.Arg(0)
		a, err = getAction(ctx, d, activity)
		if err != nil {
			return err
		}
	}

	if !useLDPreload && a.Package.Debuggable {
		ctx.Print("Package is debuggable")
	} else {
		err = d.Root(ctx)
		switch err {
		case nil:
		case adb.ErrDeviceNotRooted:
			return err
		default:
			return fmt.Errorf("Failed to restart ADB as root: %v", err)
		}
		ctx.Print("Device is rooted")
	}

	// Filenames - if no name specified, use package name.
	if output == "" {
		output = a.Package.Name + ".gfxtrace"
	}
	if inputFile == "" {
		inputFile = a.Package.Name + ".inputs"
	}

	if clearCache {
		ctx.Print("Clearing package cache")
		if err := a.Package.ClearCache(ctx); err != nil {
			return err
		}
	}

	if wasScreenOn, _ := d.IsScreenOn(ctx); !wasScreenOn {
		defer d.TurnScreenOff(ctx) // Think green!
	}

	start := gapii.JdwpStart
	if useLDPreload {
		start = gapii.AdbStart
	}

	libraries := gapii.NewLibraries()
	libraries.Set(gapii.LibInterceptor, options.Interceptor)
	libraries.Set(gapii.LibGapii, options.Gapii)
	libraries.Set(gapii.LibVkLayer, options.VkLayer)
	if err := libraries.Find(ctx, a); err != nil {
		return err
	}
	port, cleanup, err := start(ctx, libraries, a)
	if err != nil {
		return err
	}
	defer cleanup(ctx)

	ctx, stop := task.WithCancel(ctx)
	if recordInputs {
		ctx.Print("Starting input recorder")
		cleanup, err := startRecordingInputs(ctx, d, inputFile)
		if err != nil {
			return err
		}
		defer cleanup()
	} else if replayInputs {
		ctx.Print("Starting input replayer")
		if err := startReplayingInputs(ctx, d, inputFile, stop); err != nil {
			return err
		}
	}

	return capture(ctx, options, int(port), output)
}

func capture(ctx log.Context, options gapii.Options, port int, out string) error {
	ctx.Info().S("file", out).Log("Creating file")
	os.MkdirAll(filepath.Dir(out), 0755)
	file, err := os.Create(out)
	if err != nil {
		return err
	}
	defer file.Close()

	if duration == 0 {
		var cancel task.CancelFunc
		ctx, cancel = task.WithCancel(ctx)
		go func() {
			println("Press enter to stop capturing...")
			os.Stdin.Read([]byte{0})
			cancel()
		}()
	} else {
		ctx, _ = task.WithTimeout(ctx, duration)
	}
	_, err = gapii.Capture(ctx, port, file, options)
	if err != nil {
		return err
	}
	return nil
}

func getAction(ctx log.Context, d adb.Device, pattern string) (*android.Action, error) {
	re := regexp.MustCompile("(?i)" + pattern)
	packages, err := d.InstalledPackages(ctx)
	if err != nil {
		return nil, err
	}
	if len(packages) == 0 {
		return nil, fmt.Errorf("No packages found")
	}
	matchingActions := []*android.Action{}
	for _, p := range packages {
		for _, action := range p.Actions {
			if re.MatchString(action.String()) {
				matchingActions = append(matchingActions, action)
			}
		}
	}
	if len(matchingActions) == 0 {
		return nil, fmt.Errorf("No actions matching %s found", pattern)
	} else if len(matchingActions) > 1 {
		fmt.Println("Matching actions:")
		for _, test := range matchingActions {
			fmt.Print("    ")
			fmt.Println(test)
		}
		return nil, fmt.Errorf("Multiple actions matching %q found", pattern)
	}
	ctx.Info().WithValue("action", matchingActions[0]).Log("")
	return matchingActions[0], nil
}