// 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"
	"os"
	"path/filepath"
	"regexp"
	"time"

	"android.googlesource.com/platform/tools/gpu/client/adb"
	"android.googlesource.com/platform/tools/gpu/client/android"
	"android.googlesource.com/platform/tools/gpu/client/gapii"
	"android.googlesource.com/platform/tools/gpu/framework/app"
	"android.googlesource.com/platform/tools/gpu/framework/file"
	"android.googlesource.com/platform/tools/gpu/framework/log"
	"android.googlesource.com/platform/tools/gpu/framework/task"
)

var (
	device        string
	duration      time.Duration
	output        string
	localPort     int
	observeFrames uint
	observeDraws  uint
	clearCache    bool
	inputFile     string
	recordInputs  bool
	replayInputs  bool
	useLDPreload  bool
)

func init() {
	verb := &app.Verb{
		Name:      "trace",
		ShortHelp: "Captures a gfx trace from an application",
		Run:       doTrace,
	}
	verb.Flags.StringVar(&device, "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.UintVar(&observeFrames, "observeFrames", 0, "capture the framebuffer every n frames (0 to disable)")
	verb.Flags.UintVar(&observeDraws, "observeDraws", 0, "capture the framebuffer every n draws (0 to disable)")
	verb.Flags.BoolVar(&clearCache, "clearCache", 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")
	app.AddVerb(verb)
}

func doTrace(ctx log.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 && localPort == 0 {
		app.Usage(ctx, "Invalid number of arguments. Expected 1, got %d", flags.NArg())
		return nil
	}

	options := gapii.Options{
		ObserveFrameFreqency: uint32(observeFrames),
		ObserveDrawFrequency: uint32(observeDraws),
	}

	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)
}

func captureADB(ctx log.Context, flags flag.FlagSet, options gapii.Options, useLDPreload bool) error {
	activity := flags.Arg(0)
	d, err := getDevice(ctx, device)
	if err != nil {
		return err
	}

	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
	}

	binDir := file.ExecutablePath().Parent()
	port, cleanup, err := start(ctx, binDir, 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, 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 getDevice(ctx log.Context, pattern string) (adb.Device, error) {
	devices, err := adb.Devices(ctx)
	if err != nil {
		return nil, err
	}
	if len(devices) == 0 {
		return nil, fmt.Errorf("No devices found")
	}
	info := ctx.Info()
	if info.Active() {
		info.Log("Device list:")
		for _, test := range devices {
			info.S("serial", test.Info().Address.Path).Log("")
		}
	}
	matchingDevices := []adb.Device{}
	if pattern == "" {
		matchingDevices = devices
	} else {
		re := regexp.MustCompile("(?i)" + pattern)
		for _, test := range devices {
			if re.MatchString(fmt.Sprint(test)) {
				matchingDevices = append(matchingDevices, test)
			}
		}
	}
	if len(matchingDevices) == 0 {
		return nil, fmt.Errorf("No devices matching %q found", pattern)
	} else if len(matchingDevices) > 1 {
		fmt.Println("Matching devices:")
		for _, test := range matchingDevices {
			fmt.Print("    ")
			fmt.Println(test.Info().Address)
		}
		return nil, fmt.Errorf("Multiple devices matching %q found", pattern)
	}
	info.WithValue("device", matchingDevices[0]).Log("Tracing")
	return matchingDevices[0], 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
}