package main

import (
	"encoding/binary"
	"errors"
	"fmt"
	"image"
	"image/color"
	"math/bits"
)

// DecodeXWD reads an X window dump as `xwd` writes it (XWDFile.h, file version 7): the screenshot
// tool's one dependency is xwd, from the X apps the display server already ships, and the PNG is made
// here — no screenshot program, no image library.
//
// Only what an X server dumps in practice is read: a ZPixmap of 24 or 32 bits per pixel with a
// TrueColor or DirectColor visual, in either byte order. Anything else is refused by name.
func DecodeXWD(b []byte) (*image.RGBA, error) {
	const headerWords = 25
	if len(b) < headerWords*4 {
		return nil, errors.New("too short for an xwd header")
	}
	h := make([]uint32, headerWords)
	for i := range h {
		h[i] = binary.BigEndian.Uint32(b[i*4:])
	}
	headerSize, version, format := h[0], h[1], h[2]
	width, height := int(h[4]), int(h[5])
	byteOrder, bpp, bytesPerLine, class := h[7], int(h[11]), int(h[12]), h[13]
	red, green, blue := h[14], h[15], h[16]
	ncolors := int(h[19])
	switch {
	case version != 7:
		return nil, fmt.Errorf("xwd file version %d; 7 is read", version)
	case format != 2:
		return nil, fmt.Errorf("pixmap format %d; ZPixmap (2) is read", format)
	case bpp != 24 && bpp != 32:
		return nil, fmt.Errorf("%d bits per pixel; 24 and 32 are read", bpp)
	case class != 4 && class != 5:
		return nil, fmt.Errorf("visual class %d; TrueColor and DirectColor are read", class)
	case width <= 0 || height <= 0 || width > 32768 || height > 32768:
		return nil, fmt.Errorf("a %dx%d image", width, height)
	case bytesPerLine < width*bpp/8:
		return nil, fmt.Errorf("%d bytes per line for %d pixels", bytesPerLine, width)
	}
	start := int(headerSize) + ncolors*12
	if start < 0 || len(b) < start+bytesPerLine*height {
		return nil, fmt.Errorf("the image data is cut short: %d bytes, %d needed", len(b), start+bytesPerLine*height)
	}
	pixels := b[start:]
	step := bpp / 8
	channel := func(p, mask uint32) uint8 {
		if mask == 0 {
			return 0
		}
		shift := bits.TrailingZeros32(mask)
		width := bits.OnesCount32(mask)
		v := (p & mask) >> shift
		if width >= 8 {
			return uint8(v >> (width - 8))
		}
		return uint8(v << (8 - width))
	}
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		row := pixels[y*bytesPerLine:]
		for x := 0; x < width; x++ {
			px := row[x*step : x*step+step]
			var p uint32
			if byteOrder == 0 { // LSBFirst
				for i := step - 1; i >= 0; i-- {
					p = p<<8 | uint32(px[i])
				}
			} else {
				for i := 0; i < step; i++ {
					p = p<<8 | uint32(px[i])
				}
			}
			img.SetRGBA(x, y, color.RGBA{channel(p, red), channel(p, green), channel(p, blue), 0xff})
		}
	}
	return img, nil
}