package main

import (
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"image/png"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	"golang.org/x/image/webp"
)

func main() {
	quality := flag.Int("q", 95, "")
	outFmt := flag.String("out", "", "")
	workers := flag.Int("workers", runtime.NumCPU(), "")
	overwrite := flag.Bool("overwrite", false, "")
	flag.Usage = func() {}
	flag.Parse()

	files := flag.Args()
	if len(files) == 0 {
		os.Exit(1)
	}

	var forcedFmt string
	if *outFmt != "" {
		f := strings.ToLower(*outFmt)
		if f == "jpeg" {
			f = "jpg"
		}
		if f != "jpg" && f != "png" {
			log.Fatalf("unsupported --out format: %s", *outFmt)
		}
		forcedFmt = f
	}

	jobs := make(chan string)
	var wg sync.WaitGroup

	for i := 0; i < *workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for path := range jobs {
				if err := processFile(path, forcedFmt, *quality, *overwrite); err != nil {
					log.Printf("%s: %v\n", path, err)
				}
			}
		}()
	}

	for _, f := range files {
		jobs <- f
	}
	close(jobs)
	wg.Wait()
}

func processFile(path, forcedFmt string, quality int, overwrite bool) error {
	inF, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open: %w", err)
	}
	defer inF.Close()

	ext := strings.ToLower(filepath.Ext(path))
	ext = strings.TrimPrefix(ext, ".")
	var img image.Image

	if ext == "webp" {
		img, err = webp.Decode(inF)
		if err != nil {
			return fmt.Errorf("decode webp: %w", err)
		}
	} else {
		img, _, err = image.Decode(inF)
		if err != nil {
			if try, _ := inF.Seek(0, io.SeekStart); try == 0 {
				if dec, werr := webp.Decode(inF); werr == nil {
					img = dec
				} else {
					return fmt.Errorf("decode: %w (webp fallback: %v)", err, werr)
				}
			} else {
				return fmt.Errorf("decode: %w", err)
			}
		}
	}

	var target string
	if forcedFmt != "" {
		target = forcedFmt
	} else {
		switch ext {
		case "jpg", "jpeg":
			target = "jpg"
		case "png":
			target = "png"
		case "webp":
			target = "jpg"
		default:
			target = "jpg"
		}
	}

	dir := filepath.Dir(path)
	base := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
	outName := fmt.Sprintf("%s-opt.%s", base, target)
	outPath := filepath.Join(dir, outName)
	if overwrite {
		outPath = path
	}

	tmpPath := outPath + ".tmp"
	outF, err := os.Create(tmpPath)
	if err != nil {
		return fmt.Errorf("create output: %w", err)
	}

	if target == "jpg" {
		if hasAlpha(img) {
			img = removeAlpha(img)
		}
		opts := &jpeg.Options{Quality: clampQuality(quality)}
		if err := jpeg.Encode(outF, img, opts); err != nil {
			outF.Close()
			os.Remove(tmpPath)
			return fmt.Errorf("encode jpeg: %w", err)
		}
	} else if target == "png" {
		enc := png.Encoder{CompressionLevel: png.BestCompression}
		if err := enc.Encode(outF, img); err != nil {
			outF.Close()
			os.Remove(tmpPath)
			return fmt.Errorf("encode png: %w", err)
		}
	} else {
		outF.Close()
		os.Remove(tmpPath)
		return fmt.Errorf("unsupported target format: %s", target)
	}
	outF.Close()

	if overwrite {
		if err := os.Rename(tmpPath, outPath); err != nil {
			return fmt.Errorf("rename: %w", err)
		}
	} else {
		final := outPath
		i := 1
		for {
			if _, err := os.Stat(final); os.IsNotExist(err) {
				break
			}
			final = filepath.Join(dir, fmt.Sprintf("%s-opt-%d.%s", base, i, target))
			i++
		}
		if err := os.Rename(tmpPath, final); err != nil {
			return fmt.Errorf("rename: %w", err)
		}
		outPath = final
	}

	fmt.Printf("%s -> %s\n", path, outPath)
	return nil
}

func clampQuality(q int) int {
	if q < 1 {
		return 1
	}
	if q > 100 {
		return 100
	}
	return q
}

func hasAlpha(img image.Image) bool {
	switch img.(type) {
	case *image.NRGBA, *image.NRGBA64, *image.RGBA, *image.RGBA64:
		return true
	default:
		b := img.Bounds()
		if b.Empty() {
			return false
		}
		r, g, b2, a := img.At(b.Min.X, b.Min.Y).RGBA()
		return a != 0xffff || r != g || g != b2
	}
}

func removeAlpha(src image.Image) image.Image {
	b := src.Bounds()
	dst := image.NewRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
	for y := 0; y < b.Dy(); y++ {
		for x := 0; x < b.Dx(); x++ {
			dst.Set(x, y, image.White)
		}
	}
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			dst.Set(x-b.Min.X, y-b.Min.Y, src.At(x, y))
		}
	}
	return dst
}