package main

import (
	"errors"
	"flag"
	"fmt"
	"image"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/disintegration/imaging"
	"github.com/kouhin/envflag"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	etag "github.com/pablor21/echo-etag/v4"

	"image/color"
	"image/jpeg"
	_ "image/jpeg"
)

type tagSize struct {
	x int
	y int
}

var frameCount = 1337
var frameFiles = []string{}
var reSampler = imaging.Lanczos

func hashFile(fileName string) uint64 {
	h := xxhash.New64()
	f, err := os.Open(fileName)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	io.Copy(h, f)
	return h.Sum64()
}

func imageToNRGBA(input image.Image) image.NRGBA {
	out := image.NewNRGBA(input.Bounds())
	for x := 0; x < input.Bounds().Dx(); x++ {
		for y := 0; y < input.Bounds().Dy(); y++ {
			out.Set(x, y, input.At(x, y))
		}
	}
	return *out

}

func NRGBAToImage(input image.NRGBA) image.Image {
	out := image.NewRGBA(input.Bounds())
	for x := 0; x < input.Bounds().Dx(); x++ {
		for y := 0; y < input.Bounds().Dy(); y++ {
			out.Set(x, y, input.At(x, y))
		}
	}
	return out
}

func autoCropLR(input image.NRGBA) image.NRGBA {
	notOneColorLine := false
	LeftColumn := 0
	RightColumn := input.Bounds().Dy()
	for x := 0; x < input.Bounds().Dx(); x++ {
		startingColour := input.At(x, 0)
		for y := 0; y < input.Bounds().Dy(); y++ {
			if input.At(x, y) != startingColour {
				notOneColorLine = true
				log.Printf("Autocrop found first non-identical line at %d ", x)
				break
			}
		}
		if notOneColorLine {
			LeftColumn = x
			break
		}
	}
	notOneColorLine = false
	for x := input.Bounds().Dx(); x >= 0; x-- {
		startingColour := input.At(x, 0)
		for y := 0; y < input.Bounds().Dy(); y++ {
			if input.At(x, y) != startingColour {
				notOneColorLine = true
				log.Printf("Autocrop found last non-identical line at %d ", x)
				break
			}
		}
		if notOneColorLine {
			RightColumn = x
			break
		}
	}
	rect := image.Rect(LeftColumn, 0, RightColumn, input.Bounds().Dy())
	log.Printf("Autocrop used rect %d,%d,%d,%d ", LeftColumn, 0, RightColumn, input.Bounds().Dy())
	return *imaging.Crop(input.SubImage(input.Bounds()), rect)
}

func autoCropTB(input image.NRGBA) image.NRGBA {
	notOneColorLine := false
	TopColumn := 0
	BottomColumn := input.Bounds().Dx()
	for y := 0; y < input.Bounds().Dy(); y++ {
		startingColour := input.At(y, 0)
		for x := 0; x < input.Bounds().Dx(); x++ {
			if input.At(x, y) != startingColour {
				notOneColorLine = true
				log.Printf("Autocrop found first non-identical line at %d ", y)
				break
			}
		}
		if notOneColorLine {
			TopColumn = y
			break
		}
	}
	notOneColorLine = false
	for y := input.Bounds().Dy(); y >= 0; y-- {
		startingColour := input.At(y, 0)
		for x := 0; x < input.Bounds().Dx(); x++ {
			if input.At(x, y) != startingColour {
				notOneColorLine = true
				log.Printf("Autocrop found last non-identical line at %d ", y)
				break
			}
		}
		if notOneColorLine {
			BottomColumn = y
			break
		}
	}
	rect := image.Rect(0, TopColumn, input.Bounds().Dx(), BottomColumn)
	log.Printf("Autocrop used rect %d,%d,%d,%d ", 0, TopColumn, input.Bounds().Dx(), BottomColumn)
	return *imaging.Crop(input.SubImage(input.Bounds()), rect)
}

func autoCrop(input image.NRGBA, size tagSize, cropType string) image.NRGBA {
	if cropType == "LR" {
		return autoCropLR(input)
	} else if cropType == "TB" {
		return autoCropTB(input)
	} else if cropType == "requestSize" {
		if size.x > size.y {
			return autoCropLR(input)
		} else if size.x < size.y {
			return autoCropTB(input)
		}
	}
	return autoCropTB(autoCropLR(input))
}

func resizeImageFile(fileName string, size tagSize, bwBack bool, outDir string, cropType string) string {
	reader, err := os.Open(fileName)
	if err != nil {
		log.Fatal(err)
	}
	defer reader.Close()
	im, _, err := image.Decode(reader)
	if err != nil {
		log.Fatal(err)
	}
	var file *os.File
	if outDir == "%TEMP%" {
		// temp dir cannot be reliably cached, we should just continue as if we have to recreate the file.
		file, err = os.CreateTemp("", "resized") // empty string means create file in temporary directory.
		if err != nil {
			log.Fatal(err)
		}
	} else {
		outpath := filepath.Join(outDir, strconv.Itoa(size.x)+"x"+strconv.Itoa(size.y))
		if err != nil {
			log.Fatal(err)
		}
		err := os.MkdirAll(outpath, os.ModePerm)
		if err != nil {
			log.Fatal(err)
		}
		outFileName := filepath.Join(outpath, filepath.Base(fileName))
		if _, err = os.Stat(outFileName); !errors.Is(err, os.ErrNotExist) {
			// file exists (silly go has no erros.IsNot)
			log.Printf("File `%s` exists, not doing image manipulation stuff.", outFileName)
			return outFileName
		}
		file, err = os.Create(outFileName)
		if err != nil {
			log.Fatal(err)
		}
	}
	defer file.Close()
	croppedImage := autoCrop(imageToNRGBA(im), size, cropType)
	var resizedImage image.NRGBA
	if (croppedImage.Bounds().Dx() < size.x) && (croppedImage.Bounds().Dy() < size.y) {
		log.Printf("Image is smaller in both dimentions than requested (image %dx%d, requested %d,%d), scaling up", croppedImage.Bounds().Dx(), croppedImage.Bounds().Dy(), size.x, size.y)
		if croppedImage.Bounds().Dx() > croppedImage.Bounds().Dy() {
			resizedImage = *imaging.Resize(NRGBAToImage(croppedImage), size.x, 0, reSampler)
		} else {
			resizedImage = *imaging.Resize(NRGBAToImage(croppedImage), 0, size.y, reSampler)
		}
	} else {
		log.Printf("Image is bigger or perfect in at least one dimention than requested (image %dx%d, requested %d,%d), fitting (scale down with ratio preserve)", croppedImage.Bounds().Dx(), croppedImage.Bounds().Dy(), size.x, size.y)
		resizedImage = *imaging.Fit(NRGBAToImage(croppedImage), size.x, size.y, reSampler)
	}
	log.Printf("Deciding if to center (image %dx%d, requested %d,%d)", resizedImage.Bounds().Dx(), resizedImage.Bounds().Dy(), size.x, size.y)
	if (resizedImage.Bounds().Dx() < size.x) || (resizedImage.Bounds().Dy() < size.y) {
		if resizedImage.Bounds().Dx() == size.x {
			log.Printf("Centering frame in Y (%d == %d so therefore %d must be < %d)", resizedImage.Bounds().Dx(), size.x, resizedImage.Bounds().Dy(), size.y)
		} else {
			log.Printf("Centering frame in Y (%d != %d so therefore %d must be < %d)", resizedImage.Bounds().Dx(), size.x, resizedImage.Bounds().Dx(), size.x)
		}

		upLeft := image.Point{0, 0}
		lowRight := image.Point{size.x, size.y}
		emptyImage := image.NewRGBA(image.Rectangle{upLeft, lowRight})
		fillColour := resizedImage.At(0, 0)
		topLeftRed, topLeftGreen, topLeftBlue, _ := fillColour.RGBA()
		luminance, _, _ := color.RGBToYCbCr(uint8(topLeftRed), uint8(topLeftGreen), uint8(topLeftBlue))
		if bwBack {
			fillColour = color.White
			if luminance < 128 {
				fillColour = color.Black
			}
		}
		for x := 0; x < size.x; x++ {
			for y := 0; y < size.y; y++ {
				emptyImage.Set(x, y, fillColour)
				//emptyImage.Set(x, y, color.RGBA{255, 0, 0, 255})
			}
		}
		point := image.Point{0, 0}
		if resizedImage.Bounds().Dx() == size.x {
			point.Y = (size.y - resizedImage.Bounds().Dy()) / 2
		} else {
			point.X = (size.x - resizedImage.Bounds().Dx()) / 2
		}
		outputImage := imaging.Overlay(emptyImage, NRGBAToImage(resizedImage), point, 1)
		resizedImage = *outputImage
	}
	if err = jpeg.Encode(file, NRGBAToImage(resizedImage), nil); err != nil {
		log.Printf("failed to encode: %v", err)
	}
	if err != nil {
		log.Fatal(err)
	}
	return file.Name()
}

func randomNumberOfTheDay() int {
	epoch := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
	today := time.Now().Truncate(24 * time.Hour)
	daysSinceEpoch := today.Sub(epoch)
	log.Printf("Days since 2000-01-01: %d", daysSinceEpoch)
	r := rand.New(rand.NewSource(int64(daysSinceEpoch + 1)))
	return r.Intn(frameCount-1) + 1
}

func rfc2616Midnight() string {
	midnight := time.Now().Truncate(24 * time.Hour)
	return midnight.Format("Mon, 2 Jan 2006 15:04:05 MST")
}

func HeadToGetMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		if c.Request().Method == http.MethodHead {
			// Set the method to GET temporarily to reuse the handler
			c.Request().Method = http.MethodGet

			defer func() {
				c.Request().Method = http.MethodHead
			}() // Restore method after

			// Call the next handler and then clear the response body
			if err := next(c); err != nil {
				if err.Error() == echo.ErrMethodNotAllowed.Error() {
					fullResponseHeaders := c.Response().Header()
					log.Printf("headers set: %v", fullResponseHeaders)
					for k, v := range fullResponseHeaders {
						c.Response().Header().Del(k)
						log.Printf("header: %s (%d): %v", k, len(v), v)
						for _, vv := range v {
							log.Printf("value: %s", vv)
							c.Response().Header().Add(k, vv)
						}
					}
					ret := c.String(http.StatusOK, "HEAD success")
					return ret
				}

				return err
			}
		}

		return next(c)
	}
}

func OncePerDayPlease(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		c.Response().Before(func() {
			c.Response().Header().Del("Last-Modified")
			c.Response().Header().Add("Last-Modified", rfc2616Midnight())
		})

		if err := next(c); err != nil { //exec main process
			c.Error(err)
		}
		return nil
	}
}

func responder(c echo.Context, size tagSize, dataDir *string, outDir *string, bwBack *bool, cropType string) error {
	frameNumber := randomNumberOfTheDay()
	//frameNumber = rand.Intn(frameCount-1) + 1
	log.Printf("Frame number: %d", frameNumber)
	outFileName := resizeImageFile(filepath.Join(*dataDir, frameFiles[frameNumber]), size, *bwBack, *outDir, cropType)
	log.Printf("Serving: %v", outFileName)
	ret := c.File(outFileName)
	if *outDir == "%TEMP%" {
		defer os.Remove(outFileName)
	}
	hash := hashFile(outFileName)
	log.Printf("Setting ETag to: %-20d", hash)
	c.Response().Header().Del("ETag")
	c.Response().Header().Add("ETag", fmt.Sprintf("%-20d", hash))

	return ret
}

func parseSize(c echo.Context, maxWidth int, maxHeight int) tagSize {
	x, err := strconv.ParseUint(c.Param("x"), 10, 32)
	if err != nil {
		panic(err)
	}
	y, err := strconv.ParseUint(c.Param("y"), 10, 32)
	if err != nil {
		panic(err)
	}
	if int(x) > maxWidth {
		panic(errors.New(fmt.Sprintf("Requested X (%d) larger than maximum (%d)", x, maxWidth)))
	}
	if int(y) > maxHeight {
		panic(errors.New(fmt.Sprintf("Requested Y (%d) larger than maximum (%d)", y, maxHeight)))
	}
	return tagSize{int(x), int(y)}
}

func main() {
	var (
		dataDir       = flag.String("data-dir", filepath.Join("/", "var", "run", "ko", "frames"), "Data directory, where the images are")
		outDir        = flag.String("output-dir", filepath.Join(".", "output"), "Output directory (used as cache) for resized files - set to `%TEMP%` to disable caching and use temp dir, default is DataDir/output")
		bwBack        = flag.Bool("black-white-back", true, "Fill background with black or white (whichever is closest) - good for eInk, if false, uses top-left pixel of image's colour")
		listenPort    = flag.Int("listen-port", 8080, "Port to listen on")
		listenAddr    = flag.String("listen-addr", "127.0.0.1", "Address to listen on")
		defaultWidth  = flag.Int("default-width", 296, "Default width for /tag endpoint")
		defaultHeight = flag.Int("default-height", 128, "Default height for /tag endpoint")
		maxWidth      = flag.Int("max-width", 1000, "max width for /tag/x/y endpoint")
		maxHeight     = flag.Int("max-height", 1000, "max height for /tag/x/y endpoint")
		autoCropType  = flag.String("autocrop-type", "full", "Type of autocropping to apply (valid options: `full`,`LR`,`TB`,`requestSize`) - Full Autocrop, Left/Right, Top/Bottom and By Request Size, which means landscape will autocrop L/R, portrait T/B and square full")
		resizerType   = flag.String("resizer-type", "full", "Type of resampling to apply (valid options: `Lanczos`,`CatmullRom`,`MitchellNetravali`,`Linear`,`Box`,`NearestNeighbor`) https://pkg.go.dev/github.com/disintegration/imaging#ResampleFilter for details")
	)
	if err := envflag.Parse(); err != nil {
		panic(err)
	}
	var validAutoCropType = false
	for _, t := range [5]string{"full", "LR", "TB", "requestSize"} {
		if *autoCropType == t {
			validAutoCropType = true
		}
	}
	if !validAutoCropType {
		*autoCropType = "full"
	}
	switch *resizerType {
	case "Lanczos":
		reSampler = imaging.Lanczos
	case "CatmullRom":
		reSampler = imaging.CatmullRom
	case "MitchellNetravali":
		reSampler = imaging.MitchellNetravali
	case "Linear":
		reSampler = imaging.Linear
	case "Box":
		reSampler = imaging.Box
	case "NearestNeighbor":
		reSampler = imaging.NearestNeighbor
	case "BSpline":
		reSampler = imaging.BSpline
	case "Bartlett":
		reSampler = imaging.Bartlett
	case "Cosine":
		reSampler = imaging.Cosine
	case "Hamming":
		reSampler = imaging.Hamming
	case "Hann":
		reSampler = imaging.Hann
	case "Hermite":
		reSampler = imaging.Hermite
	case "Welch":
		reSampler = imaging.Welch
	default:
		reSampler = imaging.Lanczos
	}
	var defaultSize = tagSize{*defaultWidth, *defaultHeight}

	entries, err := os.ReadDir(*dataDir)
	if err != nil {
		log.Fatal(err)
	}

	for _, e := range entries {
		frameFiles = append(frameFiles, e.Name())
	}
	frameCount = len(frameFiles)
	log.Printf("Loaded %d frames from %s", frameCount, *dataDir)
	if len(frameFiles) <= 0 {
		panic("Looks like we can't find any files to display!")
	}

	e := echo.New()
	e.Use(OncePerDayPlease)
	e.Use(etag.Etag())
	e.Use(middleware.RecoverWithConfig(middleware.RecoverConfig{
		StackSize: 1 << 10, // 1 KB
		LogLevel:  1,
	}))
	e.GET("/", func(c echo.Context) error {
		return c.String(http.StatusOK, "Hello, World!")
	})
	e.GET("/tag", func(c echo.Context) error {
		return responder(c, defaultSize, dataDir, outDir, bwBack, *autoCropType)
	})
	e.HEAD("/tag", func(c echo.Context) error {
		return responder(c, defaultSize, dataDir, outDir, bwBack, *autoCropType)
	})
	e.GET("/tag/:x/:y", func(c echo.Context) error {
		size := parseSize(c, *maxWidth, *maxHeight)
		return responder(c, size, dataDir, outDir, bwBack, *autoCropType)
	})
	e.HEAD("/tag/:x/:y", func(c echo.Context) error {
		size := parseSize(c, *maxWidth, *maxHeight)
		return responder(c, size, dataDir, outDir, bwBack, *autoCropType)
	})
	e.Logger.Fatal(e.Start(*listenAddr + ":" + strconv.Itoa(*listenPort)))
}