// A VNC server.
package main

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"time"
)

type ServerInit struct {
	width, height uint16
	pixelFormat   PIXEL_FORMAT
	nameLength    uint32
}

type PIXEL_FORMAT struct {
	Bits_per_pixel, Depth, Big_endian_flag, True_color_flag byte
	Red_max, Green_max, Blue_max                            uint16
	Red_shift, Green_shift, Blue_shift                      uint8
	_, _, _                                                 byte
}

// Omits leading message-type field
type KeyEvent struct {
	DownFlag uint8
	_, _     byte // padding for encoding/binary
	KeySym   uint32
}

// Omits leading message-type field and trailing list of encoding types
type SetEncodings struct {
	_                 byte // padding for encoding/binary
	NumberOfEncodings uint16
}

type EncodingType int32

var encodingMap = map[EncodingType]string{
	0:    "Raw",
	1:    "CopyRect",
	2:    "RRE",
	4:    "CoRRE", // obsolete; present in rfbproto-3.3.pdf but not RFC6143
	5:    "Hextile",
	15:   "TRLE",
	16:   "ZRLE",
	-239: "Cursor pseudo-encoding",
	-223: "DesktopSize pseudo-encoding",
}

func (t EncodingType) String() string {
	s, ok := encodingMap[t]
	if ok {
		return s
	}
	return fmt.Sprintf("unknown encoding %d", int32(t))
}

const RawEncoding = EncodingType(0)

// Omits leading message-type field
type FramebufferUpdateRequest struct {
	Incremental                         uint8
	XPosition, YPosition, Width, Height uint16
}

// Omits leading message-type field
type PointerEvent struct {
	ButtonMask           uint8
	XPosition, YPosition uint16
}

// Omits leading message-type field
type ClientCutTextHeader struct {
	_, _, _ byte // padding for encoding/binary
	Length  uint32
}

// Includes leading message-type field
type FramebufferUpdate struct {
	MessageType        uint8
	_                  byte // padding for encoding/binary
	NumberOfRectangles uint16
}

type RectangleHeader struct {
	XPosition, YPosition, Width, Height uint16
	EncodingType                        EncodingType
}

var pixelFormat24Bit = PIXEL_FORMAT{
	Bits_per_pixel: 32, Depth: 24,
	Big_endian_flag: 1,
	True_color_flag: 1,

	Red_max: 255, Green_max: 255, Blue_max: 255,
	Red_shift: 16, Green_shift: 8, Blue_shift: 0,
}

type errConn struct {
	conn net.Conn
	err  error
	when string
}

func (c *errConn) noteErrors(when string) {
	if c.err != nil {
		c.when = when
	}
}

func (c *errConn) Write(data []byte, when string) {
	if c.err != nil {
		return
	}
	_, c.err = c.conn.Write(data)
	c.noteErrors(when)
}

func (c *errConn) Read(data []byte, when string) (count int) {
	if c.err != nil {
		return
	}
	count, c.err = c.conn.Read(data)
	c.noteErrors(when)
	return
}

func (c *errConn) ReadFull(data []byte, when string) (count int) {
	if c.err != nil {
		return
	}
	count, c.err = io.ReadFull(c.conn, data)
	c.noteErrors(when)
	return
}

func (c *errConn) WriteBigEndian(data interface{}, when string) {
	if c.err != nil {
		return
	}
	c.err = binary.Write(c.conn, binary.BigEndian, data)
	c.noteErrors(when)
}

func (c *errConn) WriteLittleEndian(data interface{}, when string) {
	if c.err != nil {
		return
	}
	c.err = binary.Write(c.conn, binary.LittleEndian, data)
	c.noteErrors(when)
}

func (c *errConn) ReadBigEndian(data interface{}, when string) {
	if c.err != nil {
		return
	}
	c.err = binary.Read(c.conn, binary.BigEndian, data)
	c.noteErrors(when)
}

func (c *errConn) Error() (error, string) {
	return c.err, c.when
}

func (c *errConn) Ok() bool {
	return c.err == nil
}

const (
	type_SetPixelFormat           = 0
	type_SetEncodings             = 2
	type_FramebufferUpdateRequest = 3
	type_KeyEvent                 = 4
	type_PointerEvent             = 5
	type_ClientCutText            = 6
)

func handleConnection(conn net.Conn) {
	defer conn.Close()
	w := &errConn{conn: conn}
	w.Write([]byte("RFB 003.003\n"), "ProtocolVersion") // ProtocolVersion message
	w.WriteBigEndian(uint32(1), "security-type")        // no authentication (security-type 1)

	// The client sends a ProtocolVersion, then \x01 for shared or
	// \x00 for noshared.
	name := "hello desktop"
	w.WriteBigEndian(ServerInit{
		width: 640, height: 480,
		pixelFormat: pixelFormat24Bit,
		nameLength:  uint32(len(name)),
	}, "ServerInit")
	w.Write([]byte(name), "name")

	data := make([]byte, 13)
	w.ReadFull(data, "client ProtocolVersion and shared/unshared")
	if w.Ok() {
		log.Printf("Successfully got ProtocolVersion and shared/unshared %q", string(data))
	}

	var (
		pixel_format  = pixelFormat24Bit
		keyEvent      KeyEvent
		setEncodings  SetEncodings
		updateRequest FramebufferUpdateRequest
		pointerEvent  PointerEvent
		cutTextHeader ClientCutTextHeader
		x, y          uint16
		start         = time.Now()
	)

	// Now read messages the client sends us and try to decode them
	for {
		err, when := w.Error()
		if err != nil {
			log.Print("closing ", conn, " because of error ", err, " in ", when)
			return
		}

		// Get message-type
		w.Read(data[:1], "message-type")
		if !w.Ok() {
			continue
		}

		message_type := data[0]
		switch message_type {
		case type_FramebufferUpdateRequest:
			w.ReadBigEndian(&updateRequest, "FramebufferUpdateRequest decoding")
			if w.Ok() {
				log.Printf("on %v got FramebufferUpdateRequest(%+v)", conn, updateRequest)

				sendFramebufferUpdate(w, &updateRequest, &pixel_format, start, x, y)
			}
		case type_PointerEvent:
			w.ReadBigEndian(&pointerEvent, "PointerEvent decoding")
			if w.Ok() {
				log.Printf("on %v got PointerEvent(%+v)", conn, pointerEvent)

				x, y = pointerEvent.XPosition, pointerEvent.YPosition
			}
		case type_KeyEvent:
			w.ReadBigEndian(&keyEvent, "KeyEvent decoding")
			if w.Ok() {
				log.Printf("on %v got KeyEvent(%#v) (%q)", conn, keyEvent, rune(keyEvent.KeySym))
			}
		case type_SetPixelFormat:
			w.ReadFull(data[:3], "SetPixelFormat padding")
			w.ReadBigEndian(&pixel_format, "SetPixelFormat decoding")
			if w.Ok() {
				log.Printf("on %v got SetPixelFormat(%+v)", conn, pixel_format)
			}
		case type_SetEncodings:
			setEncodings.NumberOfEncodings = 0 // in case of error
			w.ReadBigEndian(&setEncodings, "SetEncodings")
			encodings := make([]EncodingType, setEncodings.NumberOfEncodings)
			var n int32
			for i := 0; i < int(setEncodings.NumberOfEncodings); i++ {
				w.ReadBigEndian(&n, "an encoding type")
				encodings[i] = EncodingType(n)
			}
			if w.Ok() {
				log.Printf("on %v got SetEncodings(%+v, %+v)", conn, setEncodings, encodings)
			}
		case type_ClientCutText:
			w.ReadBigEndian(&cutTextHeader, "ClientCutText")
			if !w.Ok() {
				continue
			}
			cutTextBody := make([]byte, cutTextHeader.Length)
			w.ReadFull(cutTextBody, "ClientCutText contents")
			if w.Ok() {
				log.Printf("on %v got ClientCutText(%q)", conn, string(cutTextBody))
			}
		default:
			n := w.Read(data[:13], "reading random data")
			if w.Ok() {
				log.Printf("got message_type %d: %q on %v", message_type, string(data[:n]), conn)
			}
		}
	}
}

func sendFramebufferUpdate(c *errConn, r *FramebufferUpdateRequest, format *PIXEL_FORMAT, start time.Time, x, y uint16) {
	width, height := uint16(640), uint16(480)
	Δt := uint32(time.Now().Sub(start).Seconds() * 256)

	c.WriteBigEndian(FramebufferUpdate{
		MessageType:        0,
		NumberOfRectangles: 1,
	}, "FramebufferUpdate header")

	c.WriteBigEndian(RectangleHeader{
		XPosition: 0, YPosition: 0,
		Width: width, Height: height,
		EncodingType: RawEncoding,
	}, "Rectangle header")

	line := make([]uint32, width)
	for j := uint32(0); j < uint32(height); j++ {
		Δy := j - uint32(y)
		for i := uint32(0); i < uint32(width); i++ {
			Δx := i - uint32(x)
			line[i] = Δx*Δx + Δy*Δy - Δt
		}

		// XXX should pay attention to the *other* aspects of
		// the pixel format *too*
		if format.Big_endian_flag != 0 {
			c.WriteBigEndian(line, "sending pixel data")
		} else {
			c.WriteLittleEndian(line, "sending pixel data")
		}
	}
}

func main() {
	ln, err := net.Listen("tcp", ":5900")
	if err != nil {
		log.Fatal(err)
	}

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Print(err)
			continue
		}
		log.Print("connection ", conn, " from ", conn.RemoteAddr())
		go handleConnection(conn)
	}
}