package main

import (
	"bufio"
	"context"
	"fmt"
	_ "image/jpeg"
	"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"
)

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 DownloadImage(url, dest string) error {
	fmt.Printf("Downloading from: '%s'\n", url)

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

	// Check if the response status is OK
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("bad status: %s", resp.Status)
	}

	// 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 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 false, nil
	}

	var ok bool = true

	for _, detectedLabel := range detectLabelsResult.Labels {
		labels = append(labels, *detectedLabel.Name)
		if strings.Contains(*detectedLabel.Name, "Cat") {
			ok = false
		}
	}
	fmt.Printf("Detected: %v\n", labels)

	return ok, labels
}

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 HandleRequest(ctx context.Context, e Event) (Event, error) {
	if e.Image == "" {
		return e, nil
	}

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

	ok, _ := IsImageAcceptable(localPath)
	if !ok {
		return Event{}, nil
	}

	err = Upload(localPath, basename, e.Bucket)
	if err != nil {
		fmt.Printf("Upload failed: %v\n", err)
		return Event{}, err
	}

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

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