package aws

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsAmiLaunchPermission() *schema.Resource {
	return &schema.Resource{
		Exists: resourceAwsAmiLaunchPermissionExists,
		Create: resourceAwsAmiLaunchPermissionCreate,
		Read:   resourceAwsAmiLaunchPermissionRead,
		Delete: resourceAwsAmiLaunchPermissionDelete,

		Schema: map[string]*schema.Schema{
			"image_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"account_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
		},
	}
}

func resourceAwsAmiLaunchPermissionExists(d *schema.ResourceData, meta interface{}) (bool, error) {
	conn := meta.(*AWSClient).ec2conn

	image_id := d.Get("image_id").(string)
	account_id := d.Get("account_id").(string)
	return hasLaunchPermission(conn, image_id, account_id)
}

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

	image_id := d.Get("image_id").(string)
	account_id := d.Get("account_id").(string)

	_, err := conn.ModifyImageAttribute(&ec2.ModifyImageAttributeInput{
		ImageId:   aws.String(image_id),
		Attribute: aws.String("launchPermission"),
		LaunchPermission: &ec2.LaunchPermissionModifications{
			Add: []*ec2.LaunchPermission{
				&ec2.LaunchPermission{UserId: aws.String(account_id)},
			},
		},
	})
	if err != nil {
		return fmt.Errorf("error creating ami launch permission: %s", err)
	}

	d.SetId(fmt.Sprintf("%s-%s", image_id, account_id))
	return nil
}

func resourceAwsAmiLaunchPermissionRead(d *schema.ResourceData, meta interface{}) error {
	return nil
}

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

	image_id := d.Get("image_id").(string)
	account_id := d.Get("account_id").(string)

	_, err := conn.ModifyImageAttribute(&ec2.ModifyImageAttributeInput{
		ImageId:   aws.String(image_id),
		Attribute: aws.String("launchPermission"),
		LaunchPermission: &ec2.LaunchPermissionModifications{
			Remove: []*ec2.LaunchPermission{
				&ec2.LaunchPermission{UserId: aws.String(account_id)},
			},
		},
	})
	if err != nil {
		return fmt.Errorf("error removing ami launch permission: %s", err)
	}

	return nil
}

func hasLaunchPermission(conn *ec2.EC2, image_id string, account_id string) (bool, error) {
	attrs, err := conn.DescribeImageAttribute(&ec2.DescribeImageAttributeInput{
		ImageId:   aws.String(image_id),
		Attribute: aws.String("launchPermission"),
	})
	if err != nil {
		return false, err
	}

	for _, lp := range attrs.LaunchPermissions {
		if *lp.UserId == account_id {
			return true, nil
		}
	}
	return false, nil
}