package main

import (
	"context"
	"fmt"
	"image"
	_ "image/jpeg"
	"image/png"
	"os"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

var sess = connectAWS()

func connectAWS() *session.Session {
	sess, err := session.NewSession(
		&aws.Config{Region: aws.String("eu-central-1")},
	)

	if err != nil {
		panic(err)
	}
	return sess
}

type Event struct {
	Image  string `json:image`
	Bucket string `json:bucket`
}

func DownloadImageFromS3(bucket, key, dest string) error {
	file, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer file.Close()

	downloader := s3manager.NewDownloader(sess)
	_, err = downloader.Download(file,
		&s3.GetObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		})
	if err != nil {
		return err
	}

	return nil
}

func Upload(src, dest, bucket string) error {
	input, _ := os.Open(src)
	defer input.Close()

	uploader := s3manager.NewUploader(sess)

	_, err := uploader.Upload(&s3manager.UploadInput{
		Bucket: aws.String(bucket), // Bucket to be used
		Key:    aws.String(dest),   // Name of the file to be saved
		Body:   input,              // File
	})

	return err
}

func rgbaToGray(img image.Image) *image.Gray {
	var (
		bounds = img.Bounds()
		gray   = image.NewGray(bounds)
	)
	for x := 0; x < bounds.Max.X; x++ {
		for y := 0; y < bounds.Max.Y; y++ {
			var rgba = img.At(x, y)
			gray.Set(x, y, rgba)
		}
	}
	return gray
}

func Grayscale(srcFile, dest string) error {
	input, _ := os.Open(srcFile)
	defer input.Close()

	output, _ := os.Create(dest)
	defer output.Close()

	// Decode the image
	src, _, err := image.Decode(input)
	if err != nil {
		return err
	}

	grayscale := rgbaToGray(src)
	png.Encode(output, grayscale)

	return nil
}

func HandleRequest(ctx context.Context, e Event) (Event, error) {
	if e.Image == "" {
		return e, nil
	}

	key := e.Image
	localPath := fmt.Sprintf("/tmp/%s", key)
	err := DownloadImageFromS3(e.Bucket, e.Image, localPath)
	if err != nil {
		fmt.Printf("Could not download image: %v\n", err)
		return Event{}, err
	}

	grayscalePng := "/tmp/grayscale.png"
	err = Grayscale(localPath, grayscalePng)
	if err != nil {
		fmt.Printf("Grayscale conversion failed: %v\n", err)
		return Event{}, err
	}

	outputName := fmt.Sprintf("%s-gray.png", key)
	err = Upload(grayscalePng, outputName, e.Bucket)
	if err != nil {
		fmt.Printf("Upload failed: %v\n", err)
		return Event{}, err
	}

	return Event{outputName, e.Bucket}, nil
}

func main() {
	lambda.Start(HandleRequest)
}
