package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"slices"
	"strings"

	runewidth "github.com/mattn/go-runewidth"
	_ "golang.org/x/image/bmp"
	"golang.org/x/image/draw"
	_ "golang.org/x/image/tiff"
	_ "golang.org/x/image/webp"
	"golang.org/x/term"
)

var (
	x = flag.Int("x", 15, "set image width in columns")
	y = flag.Int("y", 0, "set image height in rows")
)

func isDir(filename string) (bool, error) {
	stat, err := os.Stat(filename)
	if err != nil {
		return false, err
	}
	return stat.IsDir(), nil
}

func isLink(filename string) (bool, error) {
	stat, err := os.Lstat(filename)
	if err != nil {
		return false, err
	}
	return stat.Mode().Type() == fs.ModeSymlink, nil
}

var errNotAnImage = errors.New("not an image")

type notAnImageError struct {
	error
}

func (e *notAnImageError) Is(err error) bool {
	return err == errNotAnImage
}

func load(filename string) (image.Image, error) {
	file := os.Stdin
	if filename != "-" {
		var err error
		file, err = os.Open(filename)
		if err != nil {
			if errors.Is(err, fs.ErrNotExist) {
				if isLink, _ := isLink(filename); isLink {
					return nil, &notAnImageError{err}
				}
			}
			return nil, err
		}
		defer file.Close()
	}

	img, _, err := image.Decode(file)
	if err != nil {
		return nil, &notAnImageError{fmt.Errorf("decode %q: %s", filename, err)}
	}

	return img, nil
}

func printImg(img image.Image, cols, lines int) (string, int) {
	// Try not to stretch the image.
	if lines <= 0 || img.Bounds().Dy()*cols <= img.Bounds().Dx()*lines*5/2 {
		lines = img.Bounds().Dy() * cols / (img.Bounds().Dx() * 5 / 2)
	} else {
		cols = img.Bounds().Dx() * lines * 5 / 2 / img.Bounds().Dy()
	}
	if lines == 0 || cols == 0 {
		return "", 0
	}

	dst := image.NewRGBA(image.Rect(0, 0, cols, 2*lines))
	draw.BiLinear.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil)

	s := new(strings.Builder)
	for y := 0; y < 2*lines; y += 2 {
		if y > 0 {
			fmt.Fprintln(s)
		}
		for x := 0; x < cols; x++ {
			hiR, hiG, hiB, _ := dst.At(x, y).RGBA()
			loR, loG, loB, _ := dst.At(x, y+1).RGBA()
			fmt.Fprintf(s, "\033[38;2;%d;%d;%dm\033[48;2;%d;%d;%dm▀",
				hiR>>8, hiG>>8, hiB>>8,
				loR>>8, loG>>8, loB>>8)
		}
		fmt.Fprint(s, "\033[39m\033[49m")
	}
	return s.String(), cols
}

type img struct {
	s        string
	width    int
	filename string
}

func compositeImages(w io.Writer, images []img, imageWidth, imagesPerRow int) {
	for row := range slices.Chunk(images, imagesPerRow) {
		rowLines := make([][]string, len(row))
		names := make([][]string, len(row))
		for i, img := range row {
			rowLines[i] = strings.Split(img.s, "\n")
			names[i] = strings.Split(runewidth.Wrap(img.filename, imageWidth), "\n")
		}
		maxLen := 0
		for imgIdx, img := range rowLines {
			maxLen = max(maxLen, len(img)+len(names[imgIdx]))
		}
		for i := 0; i < maxLen; i++ {
			if i > 0 {
				fmt.Fprintln(w)
			}
			for imgIdx, img := range rowLines {
				if imgIdx > 0 {
					fmt.Fprint(w, " ")
				}
				if i < len(img) {
					fmt.Fprintf(w, "%s%*s", img[i], max(imageWidth-row[imgIdx].width, 0), "")
				} else if i < len(img)+len(names[imgIdx]) {
					fmt.Fprintf(w, "%s", runewidth.FillRight(names[imgIdx][i-len(img)], imageWidth))
				} else {
					fmt.Fprintf(w, "%*s", imageWidth, "")
				}
			}
		}
		fmt.Fprintln(w)
	}
}

func thumbnailFiles(w io.Writer, cols, lines, imagesPerRow int, args []string) error {
	images := make([]img, len(args))
	for i, arg := range args {
		image, err := load(arg)
		if err != nil {
			return err
		}
		imageString, width := printImg(image, cols, lines)
		images[i] = img{imageString, width, filepath.Base(arg)}
	}
	compositeImages(w, images, cols, imagesPerRow)
	return nil
}

func readDir(d string) ([]string, error) {
	files, err := os.ReadDir(d)
	if err != nil {
		return nil, err
	}
	filenames := make([]string, len(files))
	for i, f := range files {
		filenames[i] = filepath.Join(d, f.Name())
	}
	return filenames, nil
}

func thumbnailDir(w io.Writer, cols, lines, imagesPerRow int, dir string) error {
	files, err := readDir(dir)
	if err != nil {
		return err
	}
	var images []img
	for _, file := range files {
		image, err := load(file)
		if err != nil {
			if errors.Is(err, errNotAnImage) {
				continue
			}
			return err
		}
		imageString, width := printImg(image, cols, lines)
		images = append(images, img{imageString, width, filepath.Base(file)})
	}
	compositeImages(w, images, cols, imagesPerRow)
	return nil
}

func thumbnailer(args []string) error {
	cols := *x
	lines := *y
	if cols <= 0 {
		return fmt.Errorf("-x must be greater than 0")
	}
	termCols, _, err := term.GetSize(1)
	if err != nil {
		return fmt.Errorf("terminal size: %s", err)
	}

	imagesPerRow := termCols / (cols + 1)

	var files, dirs []string
	if len(args) == 0 {
		dirs = []string{"."}
	} else {
		for _, arg := range args {
			isDir, err := isDir(arg)
			if err != nil {
				return err
			}
			if isDir {
				dirs = append(dirs, arg)
			} else {
				files = append(files, arg)
			}
		}
	}
	w := bufio.NewWriter(os.Stdout)
	defer w.Flush()
	if err := thumbnailFiles(w, cols, lines, imagesPerRow, files); err != nil {
		return err
	}
	if len(files) > 0 && len(dirs) > 0 {
		fmt.Fprintln(w)
	}
	for i, d := range dirs {
		if i > 0 {
			fmt.Fprintln(w)
		}
		if len(dirs) > 1 {
			fmt.Fprintf(w, "%s:\n", d)
		}
		if err := thumbnailDir(w, cols, lines, imagesPerRow, d); err != nil {
			return err
		}
	}

	return nil
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: thumbnailer [FILE]...\n")
	}
	flag.Parse()

	if err := thumbnailer(flag.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "Failed: %s\n", err)
		os.Exit(1)
	}
}