package aws

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsAmi() *schema.Resource {
	// Our schema is shared also with aws_ami_copy and aws_ami_from_instance
	resourceSchema := resourceAwsAmiCommonSchema(false)

	return &schema.Resource{
		Create: resourceAwsAmiCreate,

		Schema: resourceSchema,

		// The Read, Update and Delete operations are shared with aws_ami_copy
		// and aws_ami_from_instance, since they differ only in how the image
		// is created.
		Read:   resourceAwsAmiRead,
		Update: resourceAwsAmiUpdate,
		Delete: resourceAwsAmiDelete,
	}
}

func resourceAwsAmiCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).ec2conn

	req := &ec2.RegisterImageInput{
		Name:               aws.String(d.Get("name").(string)),
		Description:        aws.String(d.Get("description").(string)),
		Architecture:       aws.String(d.Get("architecture").(string)),
		ImageLocation:      aws.String(d.Get("image_location").(string)),
		RootDeviceName:     aws.String(d.Get("root_device_name").(string)),
		SriovNetSupport:    aws.String(d.Get("sriov_net_support").(string)),
		VirtualizationType: aws.String(d.Get("virtualization_type").(string)),
	}

	if kernelId := d.Get("kernel_id").(string); kernelId != "" {
		req.KernelId = aws.String(kernelId)
	}
	if ramdiskId := d.Get("ramdisk_id").(string); ramdiskId != "" {
		req.RamdiskId = aws.String(ramdiskId)
	}

	ebsBlockDevsSet := d.Get("ebs_block_device").(*schema.Set)
	ephemeralBlockDevsSet := d.Get("ephemeral_block_device").(*schema.Set)
	for _, ebsBlockDevI := range ebsBlockDevsSet.List() {
		ebsBlockDev := ebsBlockDevI.(map[string]interface{})
		blockDev := &ec2.BlockDeviceMapping{
			DeviceName: aws.String(ebsBlockDev["device_name"].(string)),
			Ebs: &ec2.EbsBlockDevice{
				DeleteOnTermination: aws.Bool(ebsBlockDev["delete_on_termination"].(bool)),
				VolumeSize:          aws.Int64(int64(ebsBlockDev["volume_size"].(int))),
				VolumeType:          aws.String(ebsBlockDev["volume_type"].(string)),
			},
		}
		if iops := ebsBlockDev["iops"].(int); iops != 0 {
			blockDev.Ebs.Iops = aws.Int64(int64(iops))
		}
		encrypted := ebsBlockDev["encrypted"].(bool)
		if snapshotId := ebsBlockDev["snapshot_id"].(string); snapshotId != "" {
			blockDev.Ebs.SnapshotId = aws.String(snapshotId)
			if encrypted {
				return errors.New("can't set both 'snapshot_id' and 'encrypted'")
			}
		} else if encrypted {
			blockDev.Ebs.Encrypted = aws.Bool(true)
		}
		req.BlockDeviceMappings = append(req.BlockDeviceMappings, blockDev)
	}
	for _, ephemeralBlockDevI := range ephemeralBlockDevsSet.List() {
		ephemeralBlockDev := ephemeralBlockDevI.(map[string]interface{})
		blockDev := &ec2.BlockDeviceMapping{
			DeviceName:  aws.String(ephemeralBlockDev["device_name"].(string)),
			VirtualName: aws.String(ephemeralBlockDev["virtual_name"].(string)),
		}
		req.BlockDeviceMappings = append(req.BlockDeviceMappings, blockDev)
	}

	res, err := client.RegisterImage(req)
	if err != nil {
		return err
	}

	id := *res.ImageId
	d.SetId(id)
	d.Partial(true) // make sure we record the id even if the rest of this gets interrupted
	d.Set("id", id)
	d.Set("manage_ebs_block_devices", false)
	d.SetPartial("id")
	d.SetPartial("manage_ebs_block_devices")
	d.Partial(false)

	_, err = resourceAwsAmiWaitForAvailable(id, client)
	if err != nil {
		return err
	}

	return resourceAwsAmiUpdate(d, meta)
}

func resourceAwsAmiRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).ec2conn
	id := d.Id()

	req := &ec2.DescribeImagesInput{
		ImageIds: []*string{aws.String(id)},
	}

	res, err := client.DescribeImages(req)
	if err != nil {
		return err
	}

	if len(res.Images) != 1 {
		d.SetId("")
		return nil
	}

	image := res.Images[0]
	state := *image.State

	if state == "pending" {
		// This could happen if a user manually adds an image we didn't create
		// to the state. We'll wait for the image to become available
		// before we continue. We should never take this branch in normal
		// circumstances since we would've waited for availability during
		// the "Create" step.
		image, err = resourceAwsAmiWaitForAvailable(id, client)
		if err != nil {
			return err
		}
		state = *image.State
	}

	if state == "deregistered" {
		d.SetId("")
		return nil
	}

	if state != "available" {
		return fmt.Errorf("AMI has become %s", state)
	}

	d.Set("name", image.Name)
	d.Set("description", image.Description)
	d.Set("image_location", image.ImageLocation)
	d.Set("architecture", image.Architecture)
	d.Set("kernel_id", image.KernelId)
	d.Set("ramdisk_id", image.RamdiskId)
	d.Set("root_device_name", image.RootDeviceName)
	d.Set("sriov_net_support", image.SriovNetSupport)
	d.Set("virtualization_type", image.VirtualizationType)

	var ebsBlockDevs []map[string]interface{}
	var ephemeralBlockDevs []map[string]interface{}

	for _, blockDev := range image.BlockDeviceMappings {
		if blockDev.Ebs != nil {
			ebsBlockDev := map[string]interface{}{
				"device_name":           *blockDev.DeviceName,
				"delete_on_termination": *blockDev.Ebs.DeleteOnTermination,
				"encrypted":             *blockDev.Ebs.Encrypted,
				"iops":                  0,
				"snapshot_id":           *blockDev.Ebs.SnapshotId,
				"volume_size":           int(*blockDev.Ebs.VolumeSize),
				"volume_type":           *blockDev.Ebs.VolumeType,
			}
			if blockDev.Ebs.Iops != nil {
				ebsBlockDev["iops"] = int(*blockDev.Ebs.Iops)
			}
			ebsBlockDevs = append(ebsBlockDevs, ebsBlockDev)
		} else {
			ephemeralBlockDevs = append(ephemeralBlockDevs, map[string]interface{}{
				"device_name":  *blockDev.DeviceName,
				"virtual_name": *blockDev.VirtualName,
			})
		}
	}

	d.Set("ebs_block_device", ebsBlockDevs)
	d.Set("ephemeral_block_device", ephemeralBlockDevs)

	d.Set("tags", tagsToMap(image.Tags))

	return nil
}

func resourceAwsAmiUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).ec2conn

	d.Partial(true)

	if err := setTags(client, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	if d.Get("description").(string) != "" {
		_, err := client.ModifyImageAttribute(&ec2.ModifyImageAttributeInput{
			ImageId: aws.String(d.Id()),
			Description: &ec2.AttributeValue{
				Value: aws.String(d.Get("description").(string)),
			},
		})
		if err != nil {
			return err
		}
		d.SetPartial("description")
	}

	d.Partial(false)

	return resourceAwsAmiRead(d, meta)
}

func resourceAwsAmiDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).ec2conn

	req := &ec2.DeregisterImageInput{
		ImageId: aws.String(d.Id()),
	}

	_, err := client.DeregisterImage(req)
	if err != nil {
		return err
	}

	// If we're managing the EBS snapshots then we need to delete those too.
	if d.Get("manage_ebs_snapshots").(bool) {
		errs := map[string]error{}
		ebsBlockDevsSet := d.Get("ebs_block_device").(*schema.Set)
		req := &ec2.DeleteSnapshotInput{}
		for _, ebsBlockDevI := range ebsBlockDevsSet.List() {
			ebsBlockDev := ebsBlockDevI.(map[string]interface{})
			snapshotId := ebsBlockDev["snapshot_id"].(string)
			if snapshotId != "" {
				req.SnapshotId = aws.String(snapshotId)
				_, err := client.DeleteSnapshot(req)
				if err != nil {
					errs[snapshotId] = err
				}
			}
		}

		if len(errs) > 0 {
			errParts := []string{"Errors while deleting associated EBS snapshots:"}
			for snapshotId, err := range errs {
				errParts = append(errParts, fmt.Sprintf("%s: %s", snapshotId, err))
			}
			errParts = append(errParts, "These are no longer managed by Terraform and must be deleted manually.")
			return errors.New(strings.Join(errParts, "\n"))
		}
	}

	d.SetId("")

	return nil
}

func resourceAwsAmiWaitForAvailable(id string, client *ec2.EC2) (*ec2.Image, error) {
	log.Printf("Waiting for AMI %s to become available...", id)

	req := &ec2.DescribeImagesInput{
		ImageIds: []*string{aws.String(id)},
	}
	pollsWhereNotFound := 0
	for {
		res, err := client.DescribeImages(req)
		if err != nil {
			// When using RegisterImage (for aws_ami) the AMI sometimes isn't available at all
			// right after the API responds, so we need to tolerate a couple Not Found errors
			// before an available AMI shows up.
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidAMIID.NotFound" {
				pollsWhereNotFound++
				// We arbitrarily stop polling after getting a "not found" error five times,
				// assuming that the AMI has been deleted by something other than Terraform.
				if pollsWhereNotFound > 5 {
					return nil, fmt.Errorf("gave up waiting for AMI to be created: %s", err)
				}
				time.Sleep(4 * time.Second)
				continue
			}
			return nil, fmt.Errorf("error reading AMI: %s", err)
		}

		if len(res.Images) != 1 {
			return nil, fmt.Errorf("new AMI vanished while pending")
		}

		state := *res.Images[0].State

		if state == "pending" {
			// Give it a few seconds before we poll again.
			time.Sleep(4 * time.Second)
			continue
		}

		if state == "available" {
			// We're done!
			return res.Images[0], nil
		}

		// If we're not pending or available then we're in one of the invalid/error
		// states, so stop polling and bail out.
		stateReason := *res.Images[0].StateReason
		return nil, fmt.Errorf("new AMI became %s while pending: %s", state, stateReason)
	}
}

func resourceAwsAmiCommonSchema(computed bool) map[string]*schema.Schema {
	// The "computed" parameter controls whether we're making
	// a schema for an AMI that's been implicitly registered (aws_ami_copy, aws_ami_from_instance)
	// or whether we're making a schema for an explicit registration (aws_ami).
	// When set, almost every attribute is marked as "computed".
	// When not set, only the "id" attribute is computed.
	// "name" and "description" are never computed, since they must always
	// be provided by the user.

	var virtualizationTypeDefault interface{}
	var deleteEbsOnTerminationDefault interface{}
	var sriovNetSupportDefault interface{}
	var architectureDefault interface{}
	var volumeTypeDefault interface{}
	if !computed {
		virtualizationTypeDefault = "paravirtual"
		deleteEbsOnTerminationDefault = true
		sriovNetSupportDefault = "simple"
		architectureDefault = "x86_64"
		volumeTypeDefault = "standard"
	}

	return map[string]*schema.Schema{
		"id": &schema.Schema{
			Type:     schema.TypeString,
			Computed: true,
		},
		"image_location": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: true,
			ForceNew: !computed,
		},
		"architecture": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: computed,
			ForceNew: !computed,
			Default:  architectureDefault,
		},
		"description": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
		},
		"kernel_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: computed,
			ForceNew: !computed,
		},
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
			ForceNew: true,
		},
		"ramdisk_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: computed,
			ForceNew: !computed,
		},
		"root_device_name": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: computed,
			ForceNew: !computed,
		},
		"sriov_net_support": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: computed,
			ForceNew: !computed,
			Default:  sriovNetSupportDefault,
		},
		"virtualization_type": &schema.Schema{
			Type:     schema.TypeString,
			Optional: !computed,
			Computed: computed,
			ForceNew: !computed,
			Default:  virtualizationTypeDefault,
		},

		// The following block device attributes intentionally mimick the
		// corresponding attributes on aws_instance, since they have the
		// same meaning.
		// However, we don't use root_block_device here because the constraint
		// on which root device attributes can be overridden for an instance to
		// not apply when registering an AMI.

		"ebs_block_device": &schema.Schema{
			Type:     schema.TypeSet,
			Optional: true,
			Computed: true,
			Elem: &schema.Resource{
				Schema: map[string]*schema.Schema{
					"delete_on_termination": &schema.Schema{
						Type:     schema.TypeBool,
						Optional: !computed,
						Default:  deleteEbsOnTerminationDefault,
						ForceNew: !computed,
						Computed: computed,
					},

					"device_name": &schema.Schema{
						Type:     schema.TypeString,
						Required: !computed,
						ForceNew: !computed,
						Computed: computed,
					},

					"encrypted": &schema.Schema{
						Type:     schema.TypeBool,
						Optional: !computed,
						Computed: computed,
						ForceNew: !computed,
					},

					"iops": &schema.Schema{
						Type:     schema.TypeInt,
						Optional: !computed,
						Computed: computed,
						ForceNew: !computed,
					},

					"snapshot_id": &schema.Schema{
						Type:     schema.TypeString,
						Optional: !computed,
						Computed: computed,
						ForceNew: !computed,
					},

					"volume_size": &schema.Schema{
						Type:     schema.TypeInt,
						Optional: !computed,
						Computed: true,
						ForceNew: !computed,
					},

					"volume_type": &schema.Schema{
						Type:     schema.TypeString,
						Optional: !computed,
						Computed: computed,
						ForceNew: !computed,
						Default:  volumeTypeDefault,
					},
				},
			},
			Set: func(v interface{}) int {
				var buf bytes.Buffer
				m := v.(map[string]interface{})
				buf.WriteString(fmt.Sprintf("%s-", m["device_name"].(string)))
				buf.WriteString(fmt.Sprintf("%s-", m["snapshot_id"].(string)))
				return hashcode.String(buf.String())
			},
		},

		"ephemeral_block_device": &schema.Schema{
			Type:     schema.TypeSet,
			Optional: true,
			Computed: true,
			ForceNew: true,
			Elem: &schema.Resource{
				Schema: map[string]*schema.Schema{
					"device_name": &schema.Schema{
						Type:     schema.TypeString,
						Required: !computed,
						Computed: computed,
					},

					"virtual_name": &schema.Schema{
						Type:     schema.TypeString,
						Required: !computed,
						Computed: computed,
					},
				},
			},
			Set: func(v interface{}) int {
				var buf bytes.Buffer
				m := v.(map[string]interface{})
				buf.WriteString(fmt.Sprintf("%s-", m["device_name"].(string)))
				buf.WriteString(fmt.Sprintf("%s-", m["virtual_name"].(string)))
				return hashcode.String(buf.String())
			},
		},

		"tags": tagsSchema(),

		// Not a public attribute; used to let the aws_ami_copy and aws_ami_from_instance
		// resources record that they implicitly created new EBS snapshots that we should
		// now manage. Not set by aws_ami, since the snapshots used there are presumed to
		// be independently managed.
		"manage_ebs_snapshots": &schema.Schema{
			Type:     schema.TypeBool,
			Computed: true,
			ForceNew: true,
		},
	}
}