package main

import (
	"bufio"
	"context"
	"fmt"
	"image"
	_ "image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strings"

	"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/rekognition"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"golang.org/x/image/draw"
)

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 MyEvent struct {
	Image  string `json:image`
	Bucket string `json:bucket`
}

func DownloadImage(url, dest string) error {
	// Get the data
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// Create the file
	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()

	// Write the body to file
	_, err = io.Copy(out, resp.Body)
	return err
}

func Resize(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
	}

	// Set the expected size that you want:
	dst := image.NewRGBA(image.Rect(0, 0, 400, 400))

	// Resize:
	draw.NearestNeighbor.Scale(dst, dst.Rect, src, src.Bounds(), draw.Over, nil)

	// Encode to `output`:
	png.Encode(output, dst)

	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 IsImageAcceptable(src string) (bool, []string) {
	rek := rekognition.New(sess)
	labels := make([]string, 1)

	f, err := os.Open(src)
	if err != nil {
		fmt.Println("Could not open image for recognition")
		return true, nil
	}
	reader := bufio.NewReader(f)
	content, err := ioutil.ReadAll(reader)
	if err != nil {
		fmt.Println("Could not read image for recognition")
		return true, nil
	}

	detectLabelsResult, err := rek.DetectLabels(&rekognition.DetectLabelsInput{
		Image: &rekognition.Image{
			Bytes: content,
		}})
	if err != nil {
		fmt.Printf("Could not perform image recognition: %v\n", err)
		return true, nil
	}

	var ok bool = true

	for _, detectedLabel := range detectLabelsResult.Labels {
		labels = append(labels, *detectedLabel.Name)
		if strings.Contains(*detectedLabel.Name, "Cat") {
			ok = false
		}
	}

	return ok, labels
}

func HandleRequest(ctx context.Context, e MyEvent) (string, error) {
	fmt.Printf("Got a new request: %v\n", e)
	basename := path.Base(e.Image)
	imgPath := fmt.Sprintf("/tmp/%s", basename)
	err := DownloadImage(e.Image, imgPath)
	if err != nil {
		fmt.Printf("Could not download image: %v\n", err)
		return "", err
	}

	ok, labels := IsImageAcceptable(imgPath)
	fmt.Println(labels)
	if !ok {
		return "The image cannot be published.", nil
	}

	resizedPng := "/tmp/resized.png"
	err = Resize(imgPath, resizedPng)
	if err != nil {
		fmt.Printf("Resize failed: %v\n", err)
		return "", err
	}

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

	err = Upload(resizedPng, fmt.Sprintf("%s-thumb.png", basename), e.Bucket)
	if err != nil {
		fmt.Printf("Upload failed: %v\n", err)
		return "", err
	}

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

	return fmt.Sprintf("Done: %s", basename), nil
}

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