package main

import (
	"flag"
	"fmt"
	"image-transcoder/codec"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
)

var codecs = map[string]codec.Coder{
	".jpg":  codec.JPGCoder{},
	".jpeg": codec.JPGCoder{},
	".png":  codec.PNGCoder{},
	".webp": codec.WebpCoder{},
}

var Error *log.Logger
var Warning *log.Logger

func init() {
	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
	logger := io.MultiWriter(os.Stdout)
	log.SetOutput(logger)

	Error = log.New(io.MultiWriter(os.Stderr, os.Stdout),
		fmt.Sprintf("%sERROR:%s ", "\033[0;101m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
	Warning = log.New(io.MultiWriter(os.Stdout),
		fmt.Sprintf("%sWarning:%s ", "\033[0;93m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
}

func main() {
	quality := flag.Int("quality", 80, "Quality of the image")
	to := flag.String("to", ".jpg", "Extension to transcode to")
	nosafe := flag.Bool("nosafe", false, "Prevents backup of the original before encoding")
	rm := flag.Bool("rm", false, "Removes the original after transcoding")
	w := flag.Int("w", 8, "Workers")
	flag.Parse()
	
	workers := make(chan struct{}, *w)

	if _, ok := codecs[*to]; !ok {
		Error.Println("Invalid extension")
		return
	}

	log.Printf("Using quality: %d", *quality)
	log.Printf("Transcoding to: %s", *to)
	log.Printf("Nosafe mode: %v", *nosafe)
	log.Printf("Remove mode: %v", *rm)

	var wg sync.WaitGroup
	for _, file := range flag.Args() {
		wg.Add(1)
		workers <- struct{}{}
		go func(file string) {
			defer func() { <-workers }()
			defer wg.Done()
			file = filepath.ToSlash(file)
			file = filepath.Clean(file)
			ext := filepath.Ext(file)
			coder, ok := codecs[ext]
			if !ok {
				Error.Printf("No encoder found for file: %s", file)
				return
			}

			filedir := filepath.Dir(file)

			filename := filepath.Base(file)
			if !*nosafe {
				bckp := filename + ".bak"
				bckp = filepath.Join(filedir, bckp)
				err := os.Link(file, bckp)
				if err != nil {
					Error.Printf("Failed to backup file: %v", err)
					return
				}
			}

			filehandle, err := os.Open(file)
			if err != nil {
				Error.Printf("Failed to open file: %v", err)
				return
			}
			stat, err := filehandle.Stat()
			if err != nil {
				Error.Printf("Failed to get file stats: %v", err)
				return
			}
			originalSize := stat.Size()

			log.Printf("Decoding %s", file)
			image, err := coder.Decode(filehandle)
			if err != nil {
				Error.Printf("Failed to decode file: %v", err)
				return
			}
			filehandle.Close()

			newfile := strings.Replace(filename, ext, *to, 1)
			newfile = filepath.Join(filedir, newfile)
			log.Printf("Encoding %s to %s", file, newfile)

			buf, err := coder.Encode(image, *quality)
			if err != nil {
				Error.Printf("Failed to encode image: %v", err)
				return
			}

			filehandle, err = os.Create(newfile)
			if err != nil {
				Error.Printf("Failed to create file %s: %v", newfile, err)
				return
			}

			n, err := filehandle.Write(buf.Bytes())
			if n != len(buf.Bytes()) || err != nil {
				Error.Printf("Failed to write file %s: %v", newfile, err)
				return
			}
			log.Printf("Wrote %dKB; Old size: %dKB; Compression: %.2f", n/1024, originalSize/1024, float32(n)/float32(originalSize))

			if *rm {
				if file == newfile {
					log.Printf("Won't remove newfile %s", file)
					return
				}
				log.Printf("Removing %s", file)
				err = os.Remove(file)
				if err != nil {
					Error.Printf("Failed to remove file %v: %v", file, err)
					return
				}
			}
		}(file)
	}
	wg.Wait()
}

func instrument() {
	numGoroutines := runtime.NumGoroutine()

	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	// log.Printf("%+v", m)

	sys := float64(m.Sys)
	ramUsedMB := sys / 1024 / 1024
	kbPerGoroutine := sys / 1024 / float64(numGoroutines)

	var numGoroutinesPretty string
	switch {
	case numGoroutines >= 1_000_000:
		numGoroutinesPretty = fmt.Sprintf("%.2fM", float64(numGoroutines)/1_000_000)
	case numGoroutines >= 1_000:
		numGoroutinesPretty = fmt.Sprintf("%.2fk", float64(numGoroutines)/1_000)
	default:
		numGoroutinesPretty = fmt.Sprintf("%d", numGoroutines)
	}

	fmt.Printf("\rNumber of active goroutines: %d (%s); RAM used: %.2f MB; KB per goroutine: %.2f", numGoroutines, numGoroutinesPretty, ramUsedMB, kbPerGoroutine)
}