package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/color/palette"
	_ "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/sys/unix"
)

var (
	x = flag.Int("x", 15, "set image width in columns")
	m = flagPrintMode(flag.CommandLine, "m", modeBlock24, "one of 'block', 'block24', or 'sixel'")
)

type printMode int

const (
	modeInvalid printMode = iota
	modeBlock
	modeBlock24
	modeSixel
)

type printModeValue printMode

func (m *printModeValue) String() string {
	switch printMode(*m) {
	case modeBlock:
		return "block"
	case modeBlock24:
		return "block24"
	case modeSixel:
		return "sixel"
	default:
		return "invalid"
	}
}

func (m *printModeValue) Set(s string) error {
	var mode printMode
	switch s {
	case "block":
		mode = modeBlock
	case "block24":
		mode = modeBlock24
	case "sixel":
		mode = modeSixel
	default:
		return fmt.Errorf("bad mode type %q, should be one of 'block', 'block24', or 'sixel'", s)
	}
	*m = printModeValue(mode)
	return nil
}

func flagPrintMode(fs *flag.FlagSet, name string, value printMode, usage string) *printMode {
	fs.Var((*printModeValue)(&value), name, usage)
	return &value
}

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
}

var (
	xtermBlack         = color.RGBA{R: 0x00, G: 0x00, B: 0x00, A: 0xff}
	xtermRed           = color.RGBA{R: 0xcd, G: 0x00, B: 0x00, A: 0xff}
	xtermGreen         = color.RGBA{R: 0x00, G: 0xcd, B: 0x00, A: 0xff}
	xtermYellow        = color.RGBA{R: 0xcd, G: 0xcd, B: 0x00, A: 0xff}
	xtermBlue          = color.RGBA{R: 0x00, G: 0x00, B: 0xee, A: 0xff}
	xtermMagenta       = color.RGBA{R: 0xcd, G: 0x00, B: 0xcd, A: 0xff}
	xtermCyan          = color.RGBA{R: 0x00, G: 0xcd, B: 0xcd, A: 0xff}
	xtermWhite         = color.RGBA{R: 0xe5, G: 0xe5, B: 0xe5, A: 0xff}
	xtermGray          = color.RGBA{R: 0x7f, G: 0x7f, B: 0x7f, A: 0xff}
	xtermBrightRed     = color.RGBA{R: 0xff, G: 0x00, B: 0x00, A: 0xff}
	xtermBrightGreen   = color.RGBA{R: 0x00, G: 0xff, B: 0x00, A: 0xff}
	xtermBrightYellow  = color.RGBA{R: 0xff, G: 0xff, B: 0x00, A: 0xff}
	xtermBrightBlue    = color.RGBA{R: 0x5c, G: 0x5c, B: 0xff, A: 0xff}
	xtermBrightMagenta = color.RGBA{R: 0xff, G: 0x00, B: 0xff, A: 0xff}
	xtermBrightCyan    = color.RGBA{R: 0x00, G: 0xff, B: 0xff, A: 0xff}
	xtermBrightWhite   = color.RGBA{R: 0xff, G: 0xff, B: 0xff, A: 0xff}
)

var xtermPalette = color.Palette{xtermBlack, xtermRed, xtermGreen, xtermYellow, xtermBlue, xtermMagenta, xtermCyan, xtermWhite, xtermGray, xtermBrightRed, xtermBrightGreen, xtermBrightYellow, xtermBrightBlue, xtermBrightMagenta, xtermBrightCyan, xtermBrightWhite}

func xtermColorCode(c color.Color) int {
	switch c {
	case xtermBlack:
		return 0
	case xtermRed:
		return 1
	case xtermGreen:
		return 2
	case xtermYellow:
		return 3
	case xtermBlue:
		return 4
	case xtermMagenta:
		return 5
	case xtermCyan:
		return 6
	case xtermWhite:
		return 7
	case xtermGray:
		return 60
	case xtermBrightRed:
		return 61
	case xtermBrightGreen:
		return 62
	case xtermBrightYellow:
		return 63
	case xtermBrightBlue:
		return 64
	case xtermBrightMagenta:
		return 65
	case xtermBrightCyan:
		return 66
	case xtermBrightWhite:
		return 67
	default:
		panic("not an xterm color")
	}
}

func printImgXterm(w io.Writer, img image.Image) {
	for y := 0; y < img.Bounds().Dy(); y += 2 {
		if y > 0 {
			fmt.Fprintln(w)
		}
		for x := 0; x < img.Bounds().Dx(); x++ {
			hi := img.At(x, y)
			if hi == color.Transparent {
				if x > 0 && img.At(x-1, y) != color.Transparent {
					fmt.Fprint(w, "\033[49m")
				}
				fmt.Fprint(w, " ")
			} else {
				lo := 9
				if y+1 < img.Bounds().Dy() && img.At(x, y+1) != color.Transparent {
					lo = xtermColorCode(img.At(x, y+1))
				}
				fmt.Fprintf(w, "\033[%dm\033[%dm▀", 30+xtermColorCode(hi), 40+lo)
			}
		}
		fmt.Fprint(w, "\033[39m\033[49m")
	}
}

func printImgBlock24(w io.Writer, img image.Image) {
	for y := 0; y < img.Bounds().Dy(); y += 2 {
		if y > 0 {
			fmt.Fprintln(w)
		}
		for x := 0; x < img.Bounds().Dx(); x++ {
			hi := img.At(x, y)
			if hi == color.Transparent {
				if x > 0 && img.At(x-1, y) != color.Transparent {
					fmt.Fprint(w, "\033[49m")
				}
				fmt.Fprint(w, " ")
			} else {
				if y+1 < img.Bounds().Dy() && img.At(x, y+1) != color.Transparent {
					r, g, b, _ := img.At(x, y+1).RGBA()
					fmt.Fprintf(w, "\033[48;2;%d;%d;%dm", r>>8, g>>8, b>>8)
				} else {
					fmt.Fprint(w, "\033[49m")
				}
				r, g, b, _ := hi.RGBA()
				fmt.Fprintf(w, "\033[38;2;%d;%d;%dm▀", r>>8, g>>8, b>>8)
			}
		}
		fmt.Fprint(w, "\033[39m\033[49m")
	}
}

func makePlan9ColorMap() map[color.Color]int {
	m := make(map[color.Color]int)
	for i, c := range palette.Plan9 {
		m[c] = i
	}
	return m
}

var plan9ColorMap = makePlan9ColorMap()

func scale100(c uint32) int {
	return int(c) * 100 / 0xffff
}

func printImgSixel(w io.Writer, img image.Image) {
	fmt.Fprint(w, "\033P7;1q")
	for i, c := range palette.Plan9 {
		r, g, b, _ := c.RGBA()
		fmt.Fprintf(w, "#%d;2;%d;%d;%d", i, scale100(r), scale100(g), scale100(b))
	}
	for y := 0; y < img.Bounds().Dy(); y++ {
		if y > 0 {
			if y%6 == 0 {
				fmt.Fprint(w, "-")
			} else {
				fmt.Fprint(w, "$")
			}
		}
		for x := 0; x < img.Bounds().Dx(); x++ {
			if img.At(x, y) == color.Transparent {
				fmt.Fprint(w, "?")
			} else {
				char := '?' + 1<<(y%6)
				fmt.Fprintf(w, "#%d%c", plan9ColorMap[img.At(x, y)], char)
			}
		}
	}
	fmt.Fprint(w, "\033\\")
}

func printImg(w io.Writer, img image.Image) {
	switch *m {
	case modeBlock:
		printImgXterm(w, img)
	case modeBlock24:
		printImgBlock24(w, img)
	case modeSixel:
		printImgSixel(w, img)
	default:
		panic("invalid mode")
	}
}

type pixelAspectRatio struct {
	x, y int
}

func compositeImageRow(images []image.Image, width, padding int, aspectRatio pixelAspectRatio, palette color.Palette) image.Image {
	resultWidth := width*len(images) + padding*(len(images)-1)
	maxHeight := 0
	for _, img := range images {
		maxHeight = max(maxHeight, img.Bounds().Dy()*width/(img.Bounds().Dx()*aspectRatio.y/aspectRatio.x))
	}
	bounds := image.Rect(0, 0, resultWidth, maxHeight)
	var result draw.Image
	if len(palette) > 0 {
		result = image.NewPaletted(bounds, append([]color.Color{color.Transparent}, palette...))
	} else {
		result = image.NewRGBA(bounds)
	}
	offset := 0
	for _, img := range images {
		height := img.Bounds().Dy() * width / (img.Bounds().Dx() * aspectRatio.y / aspectRatio.x)
		dstBounds := image.Rect(offset, (maxHeight-height)/2, offset+width, (maxHeight-height)/2+height)
		if len(palette) > 0 {
			tmp := image.NewRGBA(image.Rect(0, 0, width, height))
			draw.BiLinear.Scale(tmp, tmp.Bounds(), img, img.Bounds(), draw.Src, nil)
			draw.FloydSteinberg.Draw(result, dstBounds, tmp, image.Point{})
		} else {
			draw.BiLinear.Scale(result, dstBounds, img, img.Bounds(), draw.Src, nil)
		}
		offset += width + padding
	}
	return result
}

type img struct {
	i        image.Image
	filename string
}

func compositeImages(w io.Writer, images []img, width, imagesPerRow, cellX, cellY int) {
	imageWidth := width
	padding := 1
	var colorPalette color.Palette
	aspectRatio := pixelAspectRatio{1, 1}
	switch *m {
	case modeBlock:
		colorPalette = xtermPalette
		aspectRatio = pixelAspectRatio{2 * cellX, cellY}
	case modeBlock24:
		aspectRatio = pixelAspectRatio{2 * cellX, cellY}
	case modeSixel:
		imageWidth = width * cellX
		padding = cellX
		colorPalette = palette.Plan9
	}
	for row := range slices.Chunk(images, imagesPerRow) {
		images := make([]image.Image, len(row))
		for i, img := range row {
			images[i] = img.i
		}
		rowImage := compositeImageRow(images, imageWidth, padding, aspectRatio, colorPalette)
		printImg(w, rowImage)
		names := make([][]string, len(row))
		maxLen := 0
		for i, img := range row {
			names[i] = strings.Split(runewidth.Wrap(img.filename, width), "\n")
			maxLen = max(maxLen, len(names[i]))
		}
		for i := 0; i < maxLen; i++ {
			fmt.Fprintln(w)
			for imgIdx, name := range names {
				if imgIdx > 0 {
					fmt.Fprint(w, " ")
				}
				if i < len(name) {
					fmt.Fprintf(w, "%s", runewidth.FillRight(name[i], width))
				} else {
					fmt.Fprintf(w, "%*s", width, "")
				}
			}
		}
		fmt.Fprintln(w)
	}
}

func thumbnailFiles(w io.Writer, args []string, width, imagesPerRow, cellX, cellY int) error {
	images := make([]img, len(args))
	for i, arg := range args {
		image, err := load(arg)
		if err != nil {
			return err
		}
		images[i] = img{image, filepath.Base(arg)}
	}
	compositeImages(w, images, width, imagesPerRow, cellX, cellY)
	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, dir string, width, imagesPerRow, cellX, cellY int) 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
		}
		images = append(images, img{image, filepath.Base(file)})
	}
	compositeImages(w, images, width, imagesPerRow, cellX, cellY)
	return nil
}

func thumbnailer(args []string) error {
	cols := *x
	if *x <= 0 {
		return fmt.Errorf("-x must be greater than 0")
	}
	ws, err := unix.IoctlGetWinsize(int(os.Stdout.Fd()), unix.TIOCGWINSZ)
	if err != nil {
		return fmt.Errorf("terminal size: %s", err)
	}
	termCols := int(ws.Col)
	cellX, cellY := int(ws.Xpixel)/int(ws.Col), int(ws.Ypixel)/int(ws.Row)

	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, files, cols, imagesPerRow, cellX, cellY); 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, d, cols, imagesPerRow, cellX, cellY); 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)
	}
}