// 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 gapii

import (
	"fmt"
	"io"
	"net"
	"reflect"
	"time"

	"android.googlesource.com/platform/tools/gpu/client/adb"
	"android.googlesource.com/platform/tools/gpu/framework/java/jdbg"
	"android.googlesource.com/platform/tools/gpu/framework/java/jdwp"
	"android.googlesource.com/platform/tools/gpu/framework/log"
	"android.googlesource.com/platform/tools/gpu/framework/task"
)

func expect(r io.Reader, expected []byte) error {
	got := make([]byte, len(expected))
	if _, err := io.ReadFull(r, got); err != nil {
		return err
	}
	if !reflect.DeepEqual(expected, got) {
		return fmt.Errorf("Expected %v, got %v", expected, got)
	}
	return nil
}

// waitForOnCreate waits for android.app.Application.onCreate to be called, and
// then suspends the thread.
func waitForOnCreate(ctx log.Context, conn *jdwp.Connection, wakeup jdwp.ThreadID) (*jdwp.EventMethodEntry, error) {
	app, err := conn.GetClassBySignature("Landroid/app/Application;")
	if err != nil {
		return nil, err
	}

	onCreate, err := conn.GetClassMethod(app.ClassID(), "onCreate", "()V")
	if err != nil {
		return nil, err
	}

	return conn.WaitForMethodEntry(ctx, app.ClassID(), onCreate.ID, wakeup)
}

// waitForVulkanLoad for android.app.ApplicationLoaders.getClassLoader to be called,
// and then suspends the thread.
// This function is what is used to tell the vulkan loader where to search for
// layers.
func waitForVulkanLoad(ctx log.Context, conn *jdwp.Connection) (*jdwp.EventMethodEntry, error) {
	loaders, err := conn.GetClassBySignature("Landroid/app/ApplicationLoaders;")
	if err != nil {
		return nil, err
	}

	getClassLoader, err := conn.GetClassMethod(loaders.ClassID(), "getClassLoader",
		"(Ljava/lang/String;IZLjava/lang/String;Ljava/lang/String;Ljava/lang/ClassLoader;)Ljava/lang/ClassLoader;")
	if err != nil {
		return nil, err
	}

	return conn.WaitForMethodEntry(ctx, loaders.ClassID(), getClassLoader.ID, 0)
}

// loadLibrariesViaJDWP connects to the application waiting for a JDWP
// connection with the specified process id, sends a number of JDWP commands to
// load the list of libraries.
func loadLibrariesViaJDWP(ctx log.Context, libraries Libraries, pid int, filedir string, device adb.Device) error {
	const (
		reconnectAttempts = 10
		reconnectDelay    = time.Second
	)

	ctx, stop := task.WithCancel(ctx)
	defer stop()

	jdwpPort, err := adb.LocalFreeTCPPort()
	if err != nil {
		return ctx.WrapError(err, "Finding free port")
	}
	ctx = ctx.I("jdwpPort", int(jdwpPort))

	ctx.Print("Forwarding JDWP port")
	if err := device.Forward(ctx, adb.TCPPort(jdwpPort), adb.Jdwp(pid)); err != nil {
		return ctx.WrapError(err, "Setting up JDWP port forwarding")
	}
	defer device.RemoveForward(ctx, adb.TCPPort(jdwpPort))

	ctx.Print("Connecting to JDWP")

	// Create a JDWP connection with the application.
	var sock net.Conn
	var conn *jdwp.Connection
	for i := 0; i < reconnectAttempts; i++ {
		if sock, err = net.Dial("tcp", fmt.Sprintf("localhost:%v", jdwpPort)); err == nil {
			if conn, err = jdwp.Open(ctx, sock); err == nil {
				break
			}
			sock.Close()
		}
		ctx.Printf("Failed to connect: %v", err)
		time.Sleep(reconnectDelay)
	}
	if err != nil {
		return ctx.WrapError(err, "Connecting to JDWP")
	}
	defer sock.Close()

	classLoaderThread := jdwp.ThreadID(0)

	load_gapii := func(j *jdbg.JDbg, filesDir jdbg.Value) error {
		j.Class("java.io.File").New(filesDir).Call("mkdir")
		for _, library := range libraries {
			if library.Path == "" {
				continue
			}
			ctx := ctx.S("library", library.Name).I("size", int(library.Size))

			srcFile := j.Class("java.io.File").New(library.Path)
			dstFile := j.Class("java.io.File").New(filesDir, library.Name)

			srcChannel := j.Class("java.io.FileInputStream").New(srcFile).Call("getChannel")
			dstChannel := j.Class("java.io.FileOutputStream").New(dstFile).Call("getChannel")

			for remaining := library.Size; remaining > 0; {
				remaining -= dstChannel.Call("transferFrom", srcChannel, library.Size-remaining, remaining).Get().(int64)
				ctx.Printf("Copied %d bytes, %d remaining", library.Size-remaining, remaining)
			}

			srcChannel.Call("close")
			dstChannel.Call("close")

			// Load the library.
			ctx.Print("Loading library...")
			ctx.Print("Library loaded")

			// Work around for loading libraries in the N previews. See http://b/29441142.
			j.Class("java.lang.Runtime").Call("getRuntime").Call("doLoad", dstFile.Call("toString"), nil)
		}
		return nil
	}

	if len(filedir) != 0 {
		// If filedir not empty, that means we could find the filedir through normal means.
		// It also means we have a shot at tracing vulkan applications. So try that now.
		ctx.Print("Waiting for ApplicationLoaders.getClassLoader()")
		getClassLoader, err := waitForVulkanLoad(ctx, conn)

		// If err != nil that means we could not find or break in getClassLoader
		// so we have no vulkan support.
		if err == nil {
			classLoaderThread = getClassLoader.Thread
			err = jdbg.Do(conn, getClassLoader.Thread, func(j *jdbg.JDbg) error {
				newLibraryPath := j.String(":" + filedir)
				obj := j.GetStackObject("librarySearchPath").Call("concat", newLibraryPath)
				j.SetStackObject("librarySearchPath", obj)
				fileDir := j.String(filedir)
				// If successfully loaded vulkan support, then we should be good to go
				// load libgapii and friends here.
				return load_gapii(j, fileDir)
			})
			if err != nil {
				return ctx.WrapError(err, "JDWP failure")
			}
		}
	}

	// If we  did not have vulkan support, then we should try to load with
	// Application.onCreate().
	if classLoaderThread == jdwp.ThreadID(0) {
		// Wait for Application.onCreate to be called.
		ctx.Print("Waiting for Application.onCreate()")
		onCreate, err := waitForOnCreate(ctx, conn, classLoaderThread)
		if err != nil {
			return ctx.WrapError(err, "Waiting for Application.OnCreate")
		}

		// Create a JDbg session to install and load the libraries.
		ctx.Print("Installing interceptor libraries")
		err = jdbg.Do(conn, onCreate.Thread, func(j *jdbg.JDbg) error {
			// This is the /data/data/com.foo.bar/files directory where we're going to
			// place the libraries.
			filesDir := j.This().Call("getFilesDir")
			return load_gapii(j, filesDir)
		})
	}
	if err != nil {
		return ctx.WrapError(err, "JDWP failure")
	}
	return nil
}