package main

import (
	"bufio"
	"flag"
	"fmt"
	"image/jpeg"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/adrium/goheif"
	"github.com/mojitaurelie/heic2jpg/exif"
)

const (
	VERSION = "0.1.0"

	usageDescription = `heic2jpg - v%s.%s.%s.%s
Note: <file> is ignored when -dir is set

	Usage: heic2jpg [OPTIONS] <file>

OPTIONS:
`
)

type (
	filer interface {
		IsDir() bool
		Name() string
	}
)

func usage() {
	fmt.Printf(usageDescription, VERSION, runtime.Version(), runtime.GOOS, runtime.GOARCH)
	flag.PrintDefaults()
}

func version() {
	fmt.Printf("heic2jpg - v%s.%s.%s.%s\n", VERSION, runtime.Version(), runtime.GOOS, runtime.GOARCH)
	os.Exit(0)
}

func main() {
	go func() {
		if r := recover(); r != nil {
			_, err := fmt.Fprintln(os.Stderr, "E: ", r)
			if err != nil {
				os.Exit(-255)
			}
			os.Exit(-1)
		}
	}()
	flag.Usage = usage
	dir := flag.String("dir", "", "Directory to analyse")
	recursive := flag.Bool("r", false, "Analyzes the directory recursively (Require -dir)")
	jpegQuality := flag.Int("q", jpeg.DefaultQuality, "Define the jpeg quality")
	ignoreErr := flag.Bool("ignore-errors", false, "Continue to convert when the converter encounter an error")
	removeOrig := flag.Bool("remove-original", false, "Remove the heic file after conversion")
	overwriteOut := flag.Bool("overwrite", false, "Overwrite file if already exists")
	outPath := flag.String("o", "", "Output directory")
	v := flag.Bool("v", false, "Show version")
	flag.Parse()

	if *v {
		version()
	}

	if *jpegQuality < 1 || *jpegQuality > 100 {
		fmt.Println("W: Illegal value for jpeg quality, quality is set to ", jpeg.DefaultQuality)
		*jpegQuality = jpeg.DefaultQuality
	}

	var files []string
	if len(*dir) == 0 {
		if len(os.Args) == 1 {
			_, err := fmt.Fprintln(os.Stderr, "E: Missing file path")
			if err != nil {
				panic(err)
			}
			os.Exit(-1)
		}
		files = append(files, os.Args[len(os.Args)-1])
	} else {
		fmt.Printf("Analyzing folder %s\n", *dir)
		var err error
		files, err = getHEICFiles(*dir, *recursive)
		if err != nil {
			_, err := fmt.Fprintf(os.Stderr, "E: Failed to analyse directory: %v\n", err)
			if err != nil {
				panic(err)
			}
			os.Exit(-1)
		}
	}
	if len(files) == 0 {
		fmt.Println("No heic file found")
		return
	}

	fmt.Printf("%d files found with the extension .heic\n\n", len(files))

	if *removeOrig {
		fmt.Printf("WARNING: all the .heic will be deleted after the convertion\n\n")
	}

	if askForConfirmation("Do you want to start the conversion?") {
		out := *dir
		rewritePath := false
		if len(*outPath) > 0 {
			out = *outPath
			rewritePath = true
		}
		start := time.Now()
		converted := convert(out, files, *ignoreErr, *jpegQuality, rewritePath, *overwriteOut)

		if *removeOrig {
			for _, file := range converted {
				err := os.Remove(file)
				if err != nil {
					_, err := fmt.Fprintf(os.Stderr, "E: Failed to remove %s: %v\n", file, err)
					if err != nil {
						panic(err)
					}
				}
			}
		}

		fmt.Printf("%d file(s) converted in %v\n", len(converted), time.Since(start))
	}
}

func getHEICFiles(dir string, rec bool) ([]string, error) {
	var res []string
	i, err := os.Stat(dir)
	if err != nil {
		return nil, err
	}
	if !i.IsDir() {
		return nil, fmt.Errorf("%s is not a directory", dir)
	}
	if rec {
		err := filepath.Walk(dir, func(path string, info fs.FileInfo, err error) error {
			if isHEIC(info) {
				res = append(res, path)
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
		return res, nil
	}
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	for _, info := range entries {
		if isHEIC(info) {
			res = append(res, filepath.Join(dir, info.Name()))
		}
	}
	return res, nil
}

func isHEIC(f filer) bool {
	if f == nil {
		return false
	}
	return !f.IsDir() && strings.HasSuffix(f.Name(), ".heic")
}

func convert(out string, files []string, ignoreErr bool, jpegQuality int, rewritePath, overwrite bool) []string {
	var success []string
	wg := new(sync.WaitGroup)
	chunk := chunkWork(files)
	for _, files := range chunk {
		wg.Add(1)
		go func(files []string, wg *sync.WaitGroup) {
			for _, file := range files {
				// Generate new path
				fout := strings.TrimSuffix(file, ".heic") + ".jpg"
				if rewritePath {
					fout = filepath.Join(out, filepath.Base(fout))
				}

				// Ignore if file exists and overwirte arg not true
				if !overwrite {
					if _, err := os.Stat(fout); err == nil {
						continue
					}
				}

				// Open source file
				fi, err := os.Open(file)
				if err != nil {
					_, err := fmt.Fprintf(os.Stderr, "E: Failed to open %s: %v\n", file, err)
					if err != nil {
						panic(err)
					}
					if ignoreErr {
						continue
					}
					os.Exit(-1)
				}

				// Extract exif
				ex, err := goheif.ExtractExif(fi)
				if err != nil {
					fmt.Printf("W: no EXIF from %s: %v\n", file, err)
				}

				// Parse jpg to bitmap
				img, err := goheif.Decode(fi)
				if err != nil {
					err := fi.Close()
					if err != nil {
						panic(err)
					}
					_, err = fmt.Fprintf(os.Stderr, "E: Failed to decode %s: %v\n", file, err)
					if err != nil {
						panic(err)
					}
					if ignoreErr {
						continue
					}
					os.Exit(-1)
				}

				// Open out file
				fo, err := os.OpenFile(fout, os.O_RDWR|os.O_CREATE, 0644)
				if err != nil {
					err := fi.Close()
					if err != nil {
						panic(err)
					}
					_, err = fmt.Fprintf(os.Stderr, "E: Failed to create output file %s: %v\n", fout, err)
					if err != nil {
						panic(err)
					}
					if ignoreErr {
						continue
					}
					os.Exit(-1)
				}

				// Convert and write to jpg file
				w, _ := exif.NewWriterExif(fo, ex)
				err = jpeg.Encode(w, img, &jpeg.Options{
					Quality: jpegQuality,
				})
				if err != nil {
					err := fi.Close()
					if err != nil {
						panic(err)
					}
					err = fo.Close()
					if err != nil {
						panic(err)
					}
					_, err = fmt.Fprintf(os.Stderr, "E: Failed to encode %s: %v\n", fout, err)
					if err != nil {
						panic(err)
					}
					if ignoreErr {
						continue
					}
					os.Exit(-1)
				}
				err = fi.Close()
				if err != nil {
					panic(err)
				}
				err = fo.Close()
				if err != nil {
					panic(err)
				}
				success = append(success, file)
			}
			wg.Done()
		}(files, wg)
	}
	wg.Wait()
	return success
}

func chunkWork(files []string) [][]string {
	var divided [][]string
	chunkSize := (len(files) + runtime.NumCPU() - 1) / runtime.NumCPU()
	for i := 0; i < len(files); i += chunkSize {
		end := i + chunkSize

		if end > len(files) {
			end = len(files)
		}

		divided = append(divided, files[i:end])
	}
	return divided
}

func askForConfirmation(s string) bool {
	reader := bufio.NewReader(os.Stdin)
	fmt.Printf("%s [y/N]: ", s)

	response, err := reader.ReadString('\n')
	if err != nil {
		panic(err)
	}

	response = strings.ToLower(strings.TrimSpace(response))

	if response == "y" || response == "yes" {
		return true
	}
	return false
}