package main

import (
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

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

var (
	useS3    bool
	s3Bucket string
	s3Region string
	localDir string
	port     string
)

func init() {
	flag.BoolVar(&useS3, "s3", false, "Use S3")
	flag.StringVar(&s3Bucket, "bucket", "", "S3 bucket name")
	flag.StringVar(&s3Region, "region", "", "S3 region")
	flag.StringVar(&localDir, "dir", "", "Local directory")
	flag.StringVar(&port, "port", "8080", "Server port")
	flag.Parse()
}

func main() {
	if useS3 && (s3Bucket == "" || s3Region == "") {
		log.Fatal("S3 bucket and region must be provided when using S3")
	}
	if !useS3 && localDir == "" {
		log.Fatal("Local directory must be provided when not using S3")
	}

	http.HandleFunc("/api/pairs", pairsHandler)
	if !useS3 {
		http.Handle("/images/", http.StripPrefix("/images/", http.FileServer(http.Dir(localDir))))
	}
	http.Handle("/", http.FileServer(http.Dir("web"))) // Serve static files from 'web' directory
	log.Println("Server starting on port", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

func pairsHandler(w http.ResponseWriter, r *http.Request) {
	var pairs []Pair
	if useS3 {
		pairs = listS3Pairs()
	} else {
		pairs = listLocalPairs(r.Host)
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(pairs)
}

type Pair struct {
	Name  string `json:"name"`
	Left  string `json:"left"`
	Right string `json:"right"`
}

func listS3Pairs() []Pair {
	sess, err := session.NewSession(&aws.Config{
		Region: aws.String(s3Region),
	})
	if err != nil {
		log.Println("Error creating AWS session:", err)
		return nil
	}
	s3Client := s3.New(sess)
	resp, err := s3Client.ListObjectsV2(&s3.ListObjectsV2Input{
		Bucket: aws.String(s3Bucket),
	})
	if err != nil {
		log.Println("Error listing S3 objects:", err)
		return nil
	}

	pairMap := make(map[string][2]string)
	for _, obj := range resp.Contents {
		key := *obj.Key
		if strings.HasSuffix(key, "_left.jpg") {
			prefix := strings.TrimSuffix(key, "_left.jpg")
			pairMap[prefix] = [2]string{key, ""}
		} else if strings.HasSuffix(key, "_right.jpg") {
			prefix := strings.TrimSuffix(key, "_right.jpg")
			if val, ok := pairMap[prefix]; ok {
				pairMap[prefix] = [2]string{val[0], key}
			} else {
				pairMap[prefix] = [2]string{"", key}
			}
		}
	}

	var pairs []Pair
	for prefix, keys := range pairMap {
		if keys[0] != "" && keys[1] != "" {
			leftReq, _ := s3Client.GetObjectRequest(&s3.GetObjectInput{
				Bucket: aws.String(s3Bucket),
				Key:    aws.String(keys[0]),
			})
			leftURL, _ := leftReq.Presign(1 * time.Hour)
			rightReq, _ := s3Client.GetObjectRequest(&s3.GetObjectInput{
				Bucket: aws.String(s3Bucket),
				Key:    aws.String(keys[1]),
			})
			rightURL, _ := rightReq.Presign(1 * time.Hour)
			pairs = append(pairs, Pair{Name: prefix, Left: leftURL, Right: rightURL})
		}
	}
	return pairs
}

func listLocalPairs(host string) []Pair {
	files, err := os.ReadDir(localDir)
	if err != nil {
		log.Println("Error reading local directory:", err)
		return nil
	}

	pairMap := make(map[string][2]string)
	for _, file := range files {
		name := file.Name()
		if strings.HasSuffix(name, "_left.jpg") {
			prefix := strings.TrimSuffix(name, "_left.jpg")
			pairMap[prefix] = [2]string{name, ""}
		} else if strings.HasSuffix(name, "_right.jpg") {
			prefix := strings.TrimSuffix(name, "_right.jpg")
			if val, ok := pairMap[prefix]; ok {
				pairMap[prefix] = [2]string{val[0], name}
			} else {
				pairMap[prefix] = [2]string{"", name}
			}
		}
	}

	baseURL := "http://" + host + "/images/"
	var pairs []Pair
	for prefix, names := range pairMap {
		if names[0] != "" && names[1] != "" {
			pairs = append(pairs, Pair{
				Name:  prefix,
				Left:  baseURL + names[0],
				Right: baseURL + names[1],
			})
		}
	}
	return pairs
}