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

import (
	"fmt"
	"strings"

	"android.googlesource.com/platform/tools/gpu/framework/device"
	"android.googlesource.com/platform/tools/gpu/framework/log"
	"android.googlesource.com/platform/tools/gpu/framework/task"
	"android.googlesource.com/platform/tools/gpu/gapid/atom"
	"android.googlesource.com/platform/tools/gpu/gapid/atom/transform"
	"android.googlesource.com/platform/tools/gpu/gapid/config"
	"android.googlesource.com/platform/tools/gpu/gapid/database"
	"android.googlesource.com/platform/tools/gpu/gapid/gfxapi"
	"android.googlesource.com/platform/tools/gpu/gapid/gfxapi/state"
	"android.googlesource.com/platform/tools/gpu/gapid/image"
	"android.googlesource.com/platform/tools/gpu/gapid/memory"
	"android.googlesource.com/platform/tools/gpu/gapid/replay"
)

var (
	// Interface compliance tests
	_ = replay.QueryIssues(api{})
	_ = replay.QueryFramebufferAttachment(api{})
)

// makeAttachementReadable is a transformation marking all color/depth/stencil attachment
// images created via vkCreateImage atoms as readable (by patching the transfer src bit).
type makeAttachementReadable struct {
	state    *gfxapi.State
	database database.Database
}

// drawConfig is a replay.Config used by colorBufferRequest and
// depthBufferRequests.
type drawConfig struct {
}

type imgRes struct {
	img *image.Image // The image data.
	err error        // The error that occurred generating the image.
}

// framebufferRequest requests a postback of a framebuffer's attachment.
type framebufferRequest struct {
	after            atom.ID
	width, height    uint32
	attachment       gfxapi.FramebufferAttachment
	out              chan imgRes
	wireframeOverlay bool
}

// color/depth/stencil attachment bit.
func patchImageUsage(usage VkImageUsageFlags) (VkImageUsageFlags, bool) {
	hasBit := func(flag VkImageUsageFlags, bit VkImageUsageFlagBits) bool {
		return (uint32(flag) & uint32(bit)) == uint32(bit)
	}

	if hasBit(usage, VkImageUsageFlagBits_VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT) ||
		hasBit(usage, VkImageUsageFlagBits_VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT) {
		return VkImageUsageFlags(uint32(usage) | uint32(VkImageUsageFlagBits_VK_IMAGE_USAGE_TRANSFER_SRC_BIT)), true
	}
	return usage, false
}

func (t *makeAttachementReadable) Transform(ctx log.Context, id atom.ID, a atom.Atom, out atom.Writer) {
	a.Mutate(ctx, t.state, t.database, nil)
	if image, ok := a.(*VkCreateImage); ok {
		pinfo := image.PCreateInfo
		info := pinfo.Read(ctx, image, t.state, t.database, nil)

		if newUsage, changed := patchImageUsage(info.Usage); changed {
			device := image.Device
			palloc := memory.Pointer(image.PAllocator)
			pimage := memory.Pointer(image.PImage)
			result := image.Result

			info.Usage = newUsage
			newInfo := atom.Must(atom.AllocData(ctx, t.state, t.database, info))
			newAtom := NewVkCreateImage(device, newInfo.Ptr(), palloc, pimage, result)
			// Carry all non-observation extras through.
			for _, e := range image.Extras().All() {
				if _, ok := e.(*atom.Observations); !ok {
					newAtom.Extras().Add(e)
				}
			}
			// Carry observations through. We cannot merge these code with the
			// above code for handling extras together since we'd like to change
			// the observations, which are slices.
			observations := image.Extras().Observations()
			for _, r := range observations.Reads {
				// TODO: filter out the old VkImageCreateInfo. That should be done via
				// creating new observations for data we are interested from t.state.
				newAtom.AddRead(r.Range, r.ID)
			}
			// Use our new VkImageCreateInfo.
			newAtom.AddRead(newInfo.Data())
			for _, w := range observations.Writes {
				newAtom.AddWrite(w.Range, w.ID)
			}
			out.Write(ctx, id, newAtom)
			return
		}
	} else if swapchain, ok := a.(*VkCreateSwapchainKHR); ok {
		pinfo := swapchain.PCreateInfo
		info := pinfo.Read(ctx, swapchain, t.state, t.database, nil)

		if newUsage, changed := patchImageUsage(info.ImageUsage); changed {
			device := swapchain.Device
			palloc := memory.Pointer(swapchain.PAllocator)
			pswapchain := memory.Pointer(swapchain.PSwapchain)
			result := swapchain.Result

			info.ImageUsage = newUsage
			newInfo := atom.Must(atom.AllocData(ctx, t.state, t.database, info))
			newAtom := NewVkCreateSwapchainKHR(device, newInfo.Ptr(), palloc, pswapchain, result)
			for _, e := range swapchain.Extras().All() {
				if _, ok := e.(*atom.Observations); !ok {
					newAtom.Extras().Add(e)
				}
			}
			observations := swapchain.Extras().Observations()
			for _, r := range observations.Reads {
				// TODO: filter out the old VkSwapchainCreateInfoKHR. That should be done via
				// creating new observations for data we are interested from t.state.
				newAtom.AddRead(r.Range, r.ID)
			}
			newAtom.AddRead(newInfo.Data())
			for _, w := range observations.Writes {
				newAtom.AddWrite(w.Range, w.ID)
			}
			out.Write(ctx, id, newAtom)
			return
		}
	} else if createRenderPass, ok := a.(*VkCreateRenderPass); ok {
		ctx.Enter("path transformation, create render branch")
		pInfo := createRenderPass.PCreateInfo
		info := pInfo.Read(ctx, createRenderPass, t.state, t.database, nil)
		pAttachments := info.PAttachments
		attachments := pAttachments.Slice(uint64(0), uint64(info.AttachmentCount), t.state).Read(ctx, createRenderPass, t.state, t.database, nil)
		changed := false
		for i, _ := range attachments {
			if attachments[i].StoreOp == VkAttachmentStoreOp_VK_ATTACHMENT_STORE_OP_DONT_CARE {
				changed = true
				attachments[i].StoreOp = VkAttachmentStoreOp_VK_ATTACHMENT_STORE_OP_STORE
			}
		}
		// Returns if no attachment description needs to be changed
		ctx.Error().Log("after checking")
		if !changed {
			out.Write(ctx, id, a)
			return
		}
		ctx.Error().Log("changed")
		// Build new attachments data, new create info and new atom
		newAttachments := atom.Must(atom.AllocData(ctx, t.state, t.database, attachments))
		info.PAttachments = VkAttachmentDescriptionᶜᵖ(newAttachments.Ptr())
		newInfo := atom.Must(atom.AllocData(ctx, t.state, t.database, info))
		newAtom := NewVkCreateRenderPass(createRenderPass.Device,
			newInfo.Ptr(),
			memory.Pointer(createRenderPass.PAllocator),
			memory.Pointer(createRenderPass.PRenderPass),
			createRenderPass.Result)
		// Add back the extras and read/write observations
		for _, e := range createRenderPass.Extras().All() {
			if _, ok := e.(*atom.Observations); !ok {
				newAtom.Extras().Add(e)
			}
		}
		for _, r := range createRenderPass.Extras().Observations().Reads {
			newAtom.AddRead(r.Range, r.ID)
		}
		newAtom.AddRead(newInfo.Data()).AddRead(newAttachments.Data())
		for _, w := range createRenderPass.Extras().Observations().Writes {
			newAtom.AddWrite(w.Range, w.ID)
		}
		out.Write(ctx, id, newAtom)
		return
	}
	out.Write(ctx, id, a)
}

func (t *makeAttachementReadable) Flush(ctx log.Context, out atom.Writer) {}

// destroyResourceAtEOS is a transformation that destroys all active
// resources at the end of stream.
type destroyResourcesAtEOS struct {
	state    *gfxapi.State
	database database.Database
}

func (t *destroyResourcesAtEOS) Transform(ctx log.Context, id atom.ID, a atom.Atom, out atom.Writer) {
	a.Mutate(ctx, t.state, t.database, nil)
	out.Write(ctx, id, a)
}

func (t *destroyResourcesAtEOS) Flush(ctx log.Context, out atom.Writer) {
	so := getStateObject(t.state)
	id := atom.NoID
	// TODO: use the correct pAllocator once we handle it.
	p := memory.Nullptr

	// Wait all queues in all devices to finish their jobs first.
	for handle, _ := range so.Devices {
		out.Write(ctx, id, NewVkDeviceWaitIdle(handle, VkResult_VK_SUCCESS))
	}

	// Synchronization primitives.
	for handle, object := range so.Events {
		out.Write(ctx, id, NewVkDestroyEvent(object.Device, handle, p))
	}
	for handle, object := range so.Fences {
		out.Write(ctx, id, NewVkDestroyFence(object.Device, handle, p))
	}
	for handle, object := range so.Semaphores {
		out.Write(ctx, id, NewVkDestroySemaphore(object.Device, handle, p))
	}

	// Framebuffers, samplers.
	for handle, object := range so.Framebuffers {
		out.Write(ctx, id, NewVkDestroyFramebuffer(object.Device, handle, p))
	}
	for handle, object := range so.Samplers {
		out.Write(ctx, id, NewVkDestroySampler(object.Device, handle, p))
	}

	for handle, object := range so.ImageViews {
		out.Write(ctx, id, NewVkDestroyImageView(object.Device, handle, p))
	}

	// Buffers.
	for handle, object := range so.BufferViews {
		out.Write(ctx, id, NewVkDestroyBufferView(object.Device, handle, p))
	}
	for handle, object := range so.Buffers {
		out.Write(ctx, id, NewVkDestroyBuffer(object.Device, handle, p))
	}

	// Descriptor sets.
	for handle, object := range so.DescriptorPools {
		out.Write(ctx, id, NewVkDestroyDescriptorPool(object.Device, handle, p))
	}
	for handle, object := range so.DescriptorSetLayouts {
		out.Write(ctx, id, NewVkDestroyDescriptorSetLayout(object.Device, handle, p))
	}

	// Shader modules.
	for handle, object := range so.ShaderModules {
		out.Write(ctx, id, NewVkDestroyShaderModule(object.Device, handle, p))
	}

	// Pipelines.
	for handle, object := range so.Pipelines {
		out.Write(ctx, id, NewVkDestroyPipeline(object.Device, handle, p))
	}
	for handle, object := range so.PipelineLayouts {
		out.Write(ctx, id, NewVkDestroyPipelineLayout(object.Device, handle, p))
	}
	for handle, object := range so.PipelineCaches {
		out.Write(ctx, id, NewVkDestroyPipelineCache(object.Device, handle, p))
	}

	// Render passes.
	for handle, object := range so.RenderPasses {
		out.Write(ctx, id, NewVkDestroyRenderPass(object.Device, handle, p))
	}

	for handle, object := range so.QueryPools {
		out.Write(ctx, id, NewVkDestroyQueryPool(object.Device, handle, p))
	}

	// Command buffers.
	for handle, object := range so.CommandPools {
		out.Write(ctx, id, NewVkDestroyCommandPool(object.Device, handle, p))
	}

	// Swapchains.
	for handle, object := range so.Swapchains {
		out.Write(ctx, id, NewVkDestroySwapchainKHR(object.Device, handle, p))
	}

	// Memories.
	for handle, object := range so.DeviceMemories {
		out.Write(ctx, id, NewVkFreeMemory(object.Device, handle, p))
	}

	// Note: so.Images also contains Swapchain images. We do not want
	// to delete those, as that must be handled by VkDestroySwapchainKHR
	for handle, object := range so.Images {
		if !object.IsSwapchainImage {
			out.Write(ctx, id, NewVkDestroyImage(object.Device, handle, p))
		}
	}
	// Devices.
	for handle, _ := range so.Devices {
		out.Write(ctx, id, NewVkDestroyDevice(handle, p))
	}

	// Surfaces.
	for handle, object := range so.Surfaces {
		out.Write(ctx, id, NewVkDestroySurfaceKHR(object.Instance, handle, p))
	}

	// Instances.
	for handle, _ := range so.Instances {
		out.Write(ctx, id, NewVkDestroyInstance(handle, p))
	}
}

// issuesConfig is a replay.Config used by issuesRequests.
type issuesConfig struct{}

// issuesRequest requests all issues found during replay to be reported to out.
type issuesRequest struct {
	out chan<- replay.Issue
}

func (a api) Replay(
	ctx log.Context,
	intent replay.Intent,
	cfg replay.Config,
	requests []replay.Request,
	device *device.Information,
	atoms atom.List,
	out atom.Writer,
	d database.Database) error {

	transforms := atom.Transforms{}
	transforms.Add(&makeAttachementReadable{state: state.New(ctx), database: d})

	readFramebuffer := newReadFramebuffer(ctx, d)
	injector := &transform.Injector{}
	// Gathers and reports any issues found.
	var issues *findIssues

	// Terminate after all atoms of interest.
	earlyTerminator := &transform.EarlyTerminator{}

	for _, req := range requests {
		switch req := req.(type) {
		case issuesRequest:
			if issues == nil {
				issues = &findIssues{}
			}
			issues.reportTo(req.out)

		case framebufferRequest:
			earlyTerminator.Add(req.after)
			switch req.attachment {
			case gfxapi.FramebufferAttachment_Depth:
				return fmt.Errorf("Depth attachments are not currently supported")
			case gfxapi.FramebufferAttachment_Stencil:
				return fmt.Errorf("Stencil attachments are not currently supported")
			default:
				idx := uint32(req.attachment - gfxapi.FramebufferAttachment_Color0)
				readFramebuffer.Color(req.after, req.width, req.height, idx, req.out)
			}
		}
	}

	if issues != nil {
		transforms.Add(issues) // Issue reporting required.
	} else {
		transforms.Add(earlyTerminator)
	}

	// Cleanup
	transforms.Add(readFramebuffer, injector)
	transforms.Add(&destroyResourcesAtEOS{state: state.New(ctx), database: d})

	if config.DebugReplay {
		ctx.Info().Logf("Replaying %d atoms using transform chain:", len(atoms.Atoms))
		for i, t := range transforms {
			ctx.Info().Logf("(%d) %#v", i, t)
		}
	}

	if config.LogTransformsToFile {
		newTransforms := atom.Transforms{}
		newTransforms.Add(transform.NewFileLog(ctx, d, "0_original_atoms"))
		for i, t := range transforms {
			var name string
			if n, ok := t.(interface {
				Name() string
			}); ok {
				name = n.Name()
			} else {
				name = strings.Replace(fmt.Sprintf("%T", t), "*", "", -1)
			}
			newTransforms.Add(t, transform.NewFileLog(ctx, d, fmt.Sprintf("%v_atoms_after_%v", i+1, name)))
		}
		transforms = newTransforms
	}

	return catchPanics(ctx, func() { transforms.Transform(ctx, atoms, out) })
}

func catchPanics(ctx log.Context, do func()) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = ctx.AsErrorf("Panic raised: %v", e)
		}
	}()
	do()
	return nil
}

func (a api) QueryFramebufferAttachment(
	ctx log.Context,
	intent replay.Intent,
	mgr *replay.Manager,
	after atom.ID,
	width, height uint32,
	attachment gfxapi.FramebufferAttachment,
	wireframeMode replay.WireframeMode) (*image.Image, error) {

	c := drawConfig{}
	out := make(chan imgRes, 1)
	r := framebufferRequest{after: after, width: width, height: height, attachment: attachment, out: out}
	if err := mgr.Replay(ctx, intent, c, r, a); err != nil {
		return nil, err
	}
	select {
	case res := <-out:
		return res.img, res.err
	case <-task.ShouldStop(ctx):
		return nil, task.StopReason(ctx)
	}
}

func (a api) QueryIssues(
	ctx log.Context,
	intent replay.Intent,
	mgr *replay.Manager,
	out chan<- replay.Issue) {

	c := issuesConfig{}
	r := issuesRequest{out: out}
	if err := mgr.Replay(ctx, intent, c, r, a); err != nil {
		out <- replay.Issue{Atom: atom.NoID, Error: err}
		close(out)
	}
}