package main

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"path/filepath"
	"unsafe"

	"golang.org/x/sys/unix"
)

// cabletest drives the panel through drm, so /dev/fb0 holds the kernel console
// and is not worth reading. The displayed buffer is reachable from out here
// instead: GETCRTC names it, and root is allowed a handle to it without being
// drm master, so it can be read without the owning process cooperating.
const (
	drmCardGlob = "/dev/dri/card*"

	drmVblankRelative     = 0x1
	drmVblankHighCrtcMask = 0x3e
	drmVblankHighCrtcShft = 1

	xrgbRedShift   = 16
	xrgbGreenShift = 8
	xrgbBlueShift  = 0
)

func drmIOWR(nr, size uintptr) uintptr { return 3<<30 | size<<16 | 0x64<<8 | nr }

type drmModeInfo struct {
	clock                                         uint32
	hdisplay, hsyncStart, hsyncEnd, htotal, hskew uint16
	vdisplay, vsyncStart, vsyncEnd, vtotal, vscan uint16
	vrefresh, flags, typ                          uint32
	name                                          [32]byte
}

type drmModeCardRes struct {
	fbIDPtr, crtcIDPtr, connIDPtr, encIDPtr     uint64
	countFBs, countCRTCs, countConns, countEncs uint32
	minWidth, maxWidth, minHeight, maxHeight    uint32
}

type drmModeCrtc struct {
	setConnectorsPtr uint64
	countConnectors  uint32
	crtcID           uint32
	fbID             uint32
	x, y             uint32
	gammaSize        uint32
	modeValid        uint32
	mode             drmModeInfo
}

type drmModeFBCmd struct {
	fbID, width, height, pitch, bpp, depth uint32
	handle                                 uint32
}

type drmModeMapDumb struct {
	handle, pad uint32
	offset      uint64
}

type drmWaitVblank struct {
	typ      uint32
	sequence uint32
	signal   uint64
	tvSec    int64
	tvUsec   int64
}

var (
	drmGetResources = drmIOWR(0xa0, unsafe.Sizeof(drmModeCardRes{}))
	drmGetCrtc      = drmIOWR(0xa1, unsafe.Sizeof(drmModeCrtc{}))
	drmGetFB        = drmIOWR(0xad, unsafe.Sizeof(drmModeFBCmd{}))
	drmMapDumb      = drmIOWR(0xb3, unsafe.Sizeof(drmModeMapDumb{}))
	drmWaitVblankIO = drmIOWR(0x3a, 24)
)

func drmIoctl(fd int, req uintptr, arg unsafe.Pointer) error {
	if _, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(fd), req, uintptr(arg)); errno != 0 {
		return errno
	}
	return nil
}

type grabber struct {
	fd    int
	crtc  uint32
	index int
}

func openGrabber() (*grabber, error) {
	paths, err := filepath.Glob(drmCardGlob)
	if err != nil {
		return nil, err
	}
	for _, p := range paths {
		fd, err := unix.Open(p, unix.O_RDWR|unix.O_CLOEXEC, 0)
		if err != nil {
			continue
		}
		crtc, index, err := activeCrtc(fd)
		if err == nil {
			return &grabber{fd: fd, crtc: crtc, index: index}, nil
		}
		unix.Close(fd)
	}
	return nil, fmt.Errorf("no drm crtc is scanning out a framebuffer")
}

func activeCrtc(fd int) (uint32, int, error) {
	var res drmModeCardRes
	if err := drmIoctl(fd, drmGetResources, unsafe.Pointer(&res)); err != nil {
		return 0, 0, err
	}
	if res.countCRTCs == 0 {
		return 0, 0, fmt.Errorf("card has no crtcs")
	}
	crtcs := make([]uint32, res.countCRTCs)
	res.countFBs, res.countConns, res.countEncs = 0, 0, 0
	res.fbIDPtr, res.connIDPtr, res.encIDPtr = 0, 0, 0
	res.crtcIDPtr = uint64(uintptr(unsafe.Pointer(&crtcs[0])))
	if err := drmIoctl(fd, drmGetResources, unsafe.Pointer(&res)); err != nil {
		return 0, 0, err
	}
	for i, id := range crtcs {
		c := drmModeCrtc{crtcID: id}
		if err := drmIoctl(fd, drmGetCrtc, unsafe.Pointer(&c)); err != nil {
			continue
		}
		if c.fbID != 0 && c.modeValid != 0 {
			return id, i, nil
		}
	}
	return 0, 0, fmt.Errorf("no crtc is scanning out a framebuffer")
}

func (g *grabber) close() { unix.Close(g.fd) }

// Reading a buffer takes about as long as a frame, so where the read starts in
// the flip cycle decides whether it stays ahead of the writer. Woken at a
// blank, the buffer GETCRTC then names has just gone on screen.
func (g *grabber) waitVblank() error {
	v := drmWaitVblank{
		typ:      drmVblankRelative | uint32(g.index<<drmVblankHighCrtcShft)&drmVblankHighCrtcMask,
		sequence: 1,
	}
	for {
		err := drmIoctl(g.fd, drmWaitVblankIO, unsafe.Pointer(&v))
		if err == unix.EINTR {
			continue
		}
		return err
	}
}

func (g *grabber) mapFront() (mem []byte, pitch, pw, ph int, err error) {
	c := drmModeCrtc{crtcID: g.crtc}
	if err := drmIoctl(g.fd, drmGetCrtc, unsafe.Pointer(&c)); err != nil {
		return nil, 0, 0, 0, fmt.Errorf("get crtc: %w", err)
	}
	if c.fbID == 0 {
		return nil, 0, 0, 0, fmt.Errorf("crtc %d is not scanning out anything", g.crtc)
	}

	// Only a drm master or root is given a handle to someone else's framebuffer.
	fb := drmModeFBCmd{fbID: c.fbID}
	if err := drmIoctl(g.fd, drmGetFB, unsafe.Pointer(&fb)); err != nil {
		return nil, 0, 0, 0, fmt.Errorf("get fb %d: %w", c.fbID, err)
	}
	if fb.handle == 0 {
		return nil, 0, 0, 0, fmt.Errorf("kernel gave no handle for fb %d", c.fbID)
	}
	if fb.bpp != 32 {
		return nil, 0, 0, 0, fmt.Errorf("expected 32bpp, got %d", fb.bpp)
	}

	m := drmModeMapDumb{handle: fb.handle}
	if err := drmIoctl(g.fd, drmMapDumb, unsafe.Pointer(&m)); err != nil {
		return nil, 0, 0, 0, fmt.Errorf("map fb %d: %w", c.fbID, err)
	}
	pitch, pw, ph = int(fb.pitch), int(fb.width), int(fb.height)
	mem, err = unix.Mmap(g.fd, int64(m.offset), pitch*ph,
		unix.PROT_READ, unix.MAP_SHARED)
	if err != nil {
		return nil, 0, 0, 0, fmt.Errorf("mmap fb %d: %w", c.fbID, err)
	}
	return mem, pitch, pw, ph, nil
}

// Scanout memory is uncached, so a read costs a good fraction of a frame and
// could be overtaken by the next redraw. Read twice and accept only if the two
// agree: the writer leaves a buffer alone for several frames after displaying
// it, which is long enough for both. Disagreement means that reasoning is
// wrong, which is worth hearing about rather than retrying past.
func (g *grabber) frame() (*image.NRGBA, error) {
	if err := g.waitVblank(); err != nil {
		return nil, fmt.Errorf("wait for vblank: %w", err)
	}
	mem, pitch, pw, ph, err := g.mapFront()
	if err != nil {
		return nil, err
	}
	a := make([]byte, len(mem))
	copy(a, mem)
	b := make([]byte, len(mem))
	copy(b, mem)
	unix.Munmap(mem)

	if !bytes.Equal(a, b) {
		return nil, fmt.Errorf("scanout buffer was redrawn while being read")
	}
	return decode(a, pitch, pw, ph), nil
}

// The panel is landscape and every draw is turned a quarter turn on its way
// into it, so the turn is undone here to get back what a person standing in
// front of it sees.
func decode(buf []byte, pitch, pw, ph int) *image.NRGBA {
	img := image.NewNRGBA(image.Rect(0, 0, ph, pw))
	for x := 0; x < ph; x++ {
		for y := 0; y < pw; y++ {
			o := x*pitch + (pw-1-y)*4
			v := uint32(buf[o]) | uint32(buf[o+1])<<8 |
				uint32(buf[o+2])<<16 | uint32(buf[o+3])<<24
			img.SetNRGBA(x, y, color.NRGBA{
				R: uint8(v >> xrgbRedShift),
				G: uint8(v >> xrgbGreenShift),
				B: uint8(v >> xrgbBlueShift),
				A: 255,
			})
		}
	}
	return img
}