package main

import (
	"bufio"
	"flag"
	"fmt"
	"image"
	"image/color"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strconv"

	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"
	"roseh.moe/pkg/sixel"
)

var (
	x = flag.Int("x", 15, "set image width in columns")
	l = flag.Bool("l", false, "print image attributes")
	m = flagPrintMode(flag.CommandLine, "m", modeSixel, "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 {
	stat, err := os.Stat(filename)
	return err == nil && stat.IsDir()
}

func isImage(filename string) bool {
	f, err := os.Open(filename)
	if err != nil {
		return false
	}
	defer f.Close()
	_, _, err = image.DecodeConfig(f)
	return err == nil
}

type imageMetadata struct {
	format, colorModel string
	width, height      int
}

func (m *imageMetadata) String() string {
	return fmt.Sprintf("%dx%d %s %s", m.width, m.height, m.format, m.colorModel)
}

func colorModelName(m color.Model) string {
	switch m {
	case color.RGBAModel:
		return "rgba"
	case color.RGBA64Model:
		return "rgba 64 bit"
	case color.NRGBAModel:
		return "rgba not alpha premultiplied"
	case color.NRGBA64Model:
		return "rgba 64 bit not alpha premultiplied"
	case color.AlphaModel:
		return "alpha"
	case color.Alpha16Model:
		return "alpha 16 bit"
	case color.GrayModel:
		return "grayscale"
	case color.Gray16Model:
		return "grayscale 16 bit"
	case color.CMYKModel:
		return "cmyk"
	case color.NYCbCrAModel:
		return "Y'CbCr with alpha"
	case color.YCbCrModel:
		return "Y'CbCr"
	}
	if palette, ok := m.(color.Palette); ok {
		return fmt.Sprintf("palettized %d colors", len(palette))
	}
	return fmt.Sprintf("unknown %v", m)
}

func readMetadata(filename string) (imageMetadata, error) {
	f, err := os.Open(filename)
	if err != nil {
		return imageMetadata{}, err
	}
	defer f.Close()
	config, format, err := image.DecodeConfig(f)
	if err != nil {
		return imageMetadata{}, err
	}
	return imageMetadata{format: format, colorModel: colorModelName(config.ColorModel), width: config.Width, height: config.Height}, nil
}

func load(filename string) (image.Image, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

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

	return img, nil
}

func printImg(w io.Writer, img image.Image) {
	switch *m {
	case modeBlock:
		sixel.PrintXTerm16(w, img)
	case modeBlock24:
		sixel.PrintBlock(w, img)
	case modeSixel:
		sixel.Print(w, img, nil)
	default:
		panic("invalid mode")
	}
}

func escapeString(s string) string {
	for _, r := range s {
		if !strconv.IsGraphic(r) {
			return strconv.QuoteToGraphic(s)
		}
	}
	return s
}

func divRound(n, d int) int {
	return (n + d/2) / d
}

type pixelAspectRatio struct {
	x, y int
}

func layoutRow(imagePaths []string, metadata []imageMetadata, width, rowWidth, cellX int) []int {
	filenameWidths := make([]int, len(imagePaths))
	for i, path := range imagePaths {
		filenameWidths[i] = cellX * runewidth.StringWidth(escapeString(filepath.Base(path)))
		if *l {
			filenameWidths[i] = max(filenameWidths[i], cellX*runewidth.StringWidth(metadata[i].String()))
		}
	}
	var logicalWidth int
	for i, filenameWidth := range filenameWidths {
		if i > 0 {
			logicalWidth += 2 * cellX
		}
		logicalWidth += max(width, filenameWidth)
	}
	layout := make([]int, len(imagePaths))
	var offset int
	for i, filenameWidth := range filenameWidths {
		imageWidth := max(width, filenameWidth)
		layout[i] = divRound(offset*rowWidth, logicalWidth)
		offset += imageWidth + 2*cellX
	}
	return layout
}

func compositeImageRow(imagePaths []string, metadata []imageMetadata, width, rowWidth, cellX int, aspectRatio pixelAspectRatio) (image.Image, error) {
	if len(metadata) == 0 {
		metadata = make([]imageMetadata, len(imagePaths))
		for i, filename := range imagePaths {
			var err error
			if metadata[i], err = readMetadata(filename); err != nil {
				return nil, err
			}
		}
	}
	var maxHeight int
	for _, fileMetadata := range metadata {
		maxHeight = max(maxHeight, divRound(fileMetadata.height*width*aspectRatio.x, fileMetadata.width*aspectRatio.y))
	}
	layout := layoutRow(imagePaths, metadata, width, rowWidth, cellX)
	bounds := image.Rect(0, 0, rowWidth, maxHeight)
	result := image.NewRGBA(bounds)
	for i, filename := range imagePaths {
		img, err := load(filename)
		if err != nil {
			return nil, err
		}
		start := layout[i]
		end := rowWidth
		if i+1 < len(imagePaths) {
			end = layout[i+1]
		}
		height := divRound(img.Bounds().Dy()*width*aspectRatio.x, img.Bounds().Dx()*aspectRatio.y)
		offsetStart := start + (end-start-width)/2
		dstBounds := image.Rect(offsetStart, (maxHeight-height)/2, offsetStart+width, (maxHeight-height)/2+height)
		draw.BiLinear.Scale(result, dstBounds, img, img.Bounds(), draw.Src, nil)
	}
	return result, nil
}

func splitRows(imagePaths []string, metadata []imageMetadata, width, rowWidth int) [][]string {
	var rows [][]string
	for len(imagePaths) > 0 {
		row := imagePaths[:0]
		thisRowWidth := 0
		for ; len(row) < len(imagePaths); row = row[:len(row)+1] {
			nextW := max(width, runewidth.StringWidth(escapeString(filepath.Base(imagePaths[len(row)]))))
			if *l {
				nextW = max(nextW, len(metadata[len(row)].String()))
			}
			if len(row) > 0 {
				nextW += 2
			}
			if thisRowWidth > 0 && thisRowWidth+nextW > rowWidth {
				break
			}
			thisRowWidth += nextW
		}
		imagePaths = imagePaths[len(row):]
		if *l {
			metadata = metadata[len(row):]
		}
		rows = append(rows, row)
	}
	return rows
}

func thumbnailFiles(w io.Writer, imagePaths []string, width, rowWidth, cellX, cellY int) error {
	padding := 2
	aspectRatio := pixelAspectRatio{1, 1}
	switch *m {
	case modeBlock:
		aspectRatio = pixelAspectRatio{cellX, cellY}
		cellX = 1
	case modeBlock24:
		aspectRatio = pixelAspectRatio{2 * cellX, cellY}
		cellX = 1
	case modeSixel:
		padding *= cellX
	}
	var metadata []imageMetadata
	if *l {
		metadata = make([]imageMetadata, len(imagePaths))
		for i, filename := range imagePaths {
			var err error
			if metadata[i], err = readMetadata(filename); err != nil {
				return err
			}
		}
	}
	rows := splitRows(imagePaths, metadata, width, rowWidth)
	rowMetadata := make([][]imageMetadata, len(rows))
	if *l {
		for i, row := range rows {
			rowMetadata[i] = metadata[:len(row)]
			metadata = metadata[len(row):]
		}
	}
	rowImages := make([]chan image.Image, len(rows))
	for i := range rowImages {
		rowImages[i] = make(chan image.Image, 1)
	}
	errCh := make(chan error, 1)
	go func() {
		sem := make(chan struct{}, 20)
		for i, row := range rows {
			sem <- struct{}{}
			go func() {
				defer func() { <-sem }()
				img, err := compositeImageRow(row, rowMetadata[i], width*cellX, rowWidth*cellX, cellX, aspectRatio)
				if err != nil {
					select {
					case errCh <- err:
					default:
					}
					return
				}
				rowImages[i] <- img
			}()
		}
	}()
	for i, row := range rows {
		var rowImage image.Image
		select {
		case err := <-errCh:
			return err
		case rowImage = <-rowImages[i]:
		}
		printImg(w, rowImage)
		fmt.Fprintln(w)
		layouts := layoutRow(row, rowMetadata[i], width*cellX, rowWidth*cellX, cellX)
		realOffset := 0
		for j, filename := range row {
			start := divRound(layouts[j], cellX)
			end := rowWidth
			if j+1 < len(row) {
				end = divRound(layouts[j+1], cellX)
			}
			fmt.Fprintf(w, "%*s", start-realOffset, "")
			realOffset = start
			fileNameWidth := runewidth.StringWidth(escapeString(filepath.Base(filename)))
			padding := max((end-start-fileNameWidth)/2, 0)
			fmt.Fprintf(w, "%*s%s", padding, "", escapeString(filepath.Base(filename)))
			realOffset += padding + fileNameWidth
		}
		fmt.Fprintln(w)
		if *l {
			realOffset = 0
			for j, metadata := range rowMetadata[i] {
				start := divRound(layouts[j], cellX)
				end := rowWidth
				if j+1 < len(row) {
					end = divRound(layouts[j+1], cellX)
				}
				fmt.Fprintf(w, "%*s", start-realOffset, "")
				realOffset = start
				details := metadata.String()
				padding := max((end-start-len(details))/2, 0)
				fmt.Fprintf(w, "%*s%s", padding, "", details)
				realOffset += padding + len(details)
			}
			fmt.Fprintln(w)
		}
	}
	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, rowWidth, cellX, cellY int) error {
	files, err := readDir(dir)
	if err != nil {
		return err
	}
	files = slices.DeleteFunc(files, func(filename string) bool { return !isImage(filename) })
	return thumbnailFiles(w, files, width, rowWidth, cellX, cellY)
}

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)

	cols = min(cols, termCols)

	var files, dirs []string
	if len(args) == 0 {
		dirs = []string{"."}
	} else {
		for _, arg := range args {
			if isDir(arg) {
				dirs = append(dirs, arg)
			} else {
				files = append(files, arg)
			}
		}
	}
	w := bufio.NewWriter(os.Stdout)
	defer w.Flush()
	if err := thumbnailFiles(w, files, cols, termCols, 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, termCols, cellX, cellY); err != nil {
			return err
		}
	}

	return nil
}

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

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