package rdp

import (
	"bytes"
	"encoding/base64"
	"encoding/binary"
	"image"
	"image/png"
	"io"

	"ezops/pkg/rdp/rle"
	"ezops/pkg/rdp/trace"
)

const (
	bitmapCompression         = 0x0001
	noBitmapCompressionHdr    = 0x0400
	bitmapNoProcessing        = 0x8000
	fastPathUpdateBitmap      = 0x1
	fastPathUpdateSurfaceCmds = 0x4
	pduType2Update            = 0x02
)

// BitmapRect 表示一块待绘制到 Canvas 的位图区域 // Logan.Li
type BitmapRect struct {
	DestLeft     int
	DestTop      int
	Width        int
	Height       int
	BitsPerPixel int
	Data         []byte
}

func bppFromWire(v uint16) int {
	switch v {
	case 8:
		return 1
	case 15, 16:
		return 2
	case 24:
		return 3
	case 32:
		return 4
	default:
		if v <= 4 {
			return int(v)
		}
		return 4
	}
}

type wireBitmap struct {
	DestLeft         uint16
	DestTop          uint16
	DestRight        uint16
	DestBottom       uint16
	Width            uint16
	Height           uint16
	BitsPerPixel     uint16
	Flags            uint16
	BitmapLength     uint16
	BitmapDataStream []byte
}

func readWireBitmap(r *bytes.Reader) (wireBitmap, error) {
	var rect wireBitmap
	var err error
	rect.DestLeft, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.DestTop, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.DestRight, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.DestBottom, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.Width, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.Height, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.BitsPerPixel, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.Flags, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	rect.BitmapLength, err = readLE16(r)
	if err != nil {
		return rect, err
	}
	ln := int(rect.BitmapLength)
	if rect.Flags&bitmapCompression != 0 && rect.Flags&noBitmapCompressionHdr == 0 {
		if _, err = readLE16(r); err != nil {
			return rect, err
		}
		mainBody, err := readLE16(r)
		if err != nil {
			return rect, err
		}
		if _, err = readLE16(r); err != nil {
			return rect, err
		}
		if _, err = readLE16(r); err != nil {
			return rect, err
		}
		ln = int(mainBody)
	}
	rect.BitmapDataStream = make([]byte, ln)
	_, err = io.ReadFull(r, rect.BitmapDataStream)
	return rect, err
}

func parseFastPathBitmap(payload []byte) ([]BitmapRect, error) {
	r := bytes.NewReader(payload)
	if _, err := readLE16(r); err != nil {
		return nil, err
	}
	count, err := readLE16(r)
	if err != nil {
		return nil, err
	}
	out := make([]BitmapRect, 0, count)
	for i := 0; i < int(count); i++ {
		wb, err := readWireBitmap(r)
		if err != nil {
			return out, err
		}
		if rect, ok := decodeWireBitmap(wb); ok {
			out = append(out, rect)
		}
	}
	return out, nil
}

func parseSlowPathBitmap(payload []byte) ([]BitmapRect, error) {
	if len(payload) < 15 {
		return nil, nil
	}
	pduType := binaryLittleUint16(payload[2:4])
	if pduType != pduTypeData {
		return nil, nil
	}
	type2 := payload[14]
	if type2 != pduType2Update {
		return nil, nil
	}
	body := payload[18:]
	if len(body) < 2 {
		return nil, nil
	}
	updateType := binaryLittleUint16(body[0:2])
	if updateType != fastPathUpdateBitmap {
		return nil, nil
	}
	return parseFastPathBitmap(body[2:])
}

func decodeWireBitmap(wb wireBitmap) (rect BitmapRect, ok bool) {
	defer func() {
		if recover() != nil {
			rect = BitmapRect{}
			ok = false
		}
	}()
	bpp := bppFromWire(wb.BitsPerPixel)
	data := wb.BitmapDataStream
	originalLen := len(data)
	if wb.Flags&bitmapNoProcessing != 0 {
		// Surface command 已是自上而下 BGRA // Logan.Li
	} else if wb.Flags&bitmapCompression != 0 {
		data = rle.DecompressInto(data, nil, int(wb.Width), int(wb.Height), bpp)
		if len(data) == 0 {
			trace.Logf("bitmap decode failed compressed rect=%dx%d@%d,%d wireBpp=%d bpp=%d flags=0x%04x original=%d", wb.Width, wb.Height, wb.DestLeft, wb.DestTop, wb.BitsPerPixel, bpp, wb.Flags, originalLen)
			return BitmapRect{}, false
		}
	} else {
		data = flipBitmapVertical(data, int(wb.Width), bpp, int(wb.Height))
	}
	if len(data) < int(wb.Width)*int(wb.Height)*bpp {
		trace.Logf("bitmap decode short rect=%dx%d@%d,%d wireBpp=%d bpp=%d flags=0x%04x original=%d decoded=%d", wb.Width, wb.Height, wb.DestLeft, wb.DestTop, wb.BitsPerPixel, bpp, wb.Flags, originalLen, len(data))
		return BitmapRect{}, false
	}
	if isMostlyBlack(data, bpp) && int(wb.Width) >= 128 && int(wb.Height) >= 128 {
		trace.Logf("bitmap decode mostly black rect=%dx%d@%d,%d wireBpp=%d bpp=%d flags=0x%04x original=%d decoded=%d", wb.Width, wb.Height, wb.DestLeft, wb.DestTop, wb.BitsPerPixel, bpp, wb.Flags, originalLen, len(data))
	}
	return BitmapRect{
		DestLeft:     int(wb.DestLeft),
		DestTop:      int(wb.DestTop),
		Width:        int(wb.Width),
		Height:       int(wb.Height),
		BitsPerPixel: bpp,
		Data:         data,
	}, true
}

func isMostlyBlack(data []byte, bpp int) bool {
	if len(data) == 0 || bpp <= 0 {
		return false
	}
	limit := len(data)
	if limit > 256*1024 {
		limit = 256 * 1024
	}
	nonBlack := 0
	for i := 0; i+bpp-1 < limit; i += bpp {
		switch bpp {
		case 1:
			if data[i] > 4 {
				nonBlack++
			}
		case 2:
			if data[i] > 0 || data[i+1] > 0 {
				nonBlack++
			}
		default:
			if data[i] > 4 || data[i+1] > 4 || data[i+2] > 4 {
				nonBlack++
			}
		}
		if nonBlack > 32 {
			return false
		}
	}
	return true
}

func flipBitmapVertical(data []byte, width, bpp, height int) []byte {
	out := append([]byte(nil), data...)
	stride := width * bpp
	tmp := make([]byte, stride)
	for y := 0; y < height/2; y++ {
		top := y * stride
		bot := (height - 1 - y) * stride
		copy(tmp, out[top:top+stride])
		copy(out[top:top+stride], out[bot:bot+stride])
		copy(out[bot:bot+stride], tmp)
	}
	return out
}

func bitmapRectPNG(rect BitmapRect) (string, error) {
	img := image.NewRGBA(image.Rect(0, 0, rect.Width, rect.Height))
	pix := img.Pix
	data := rect.Data
	switch rect.BitsPerPixel {
	case 1, 2:
		n := rect.Width * rect.Height
		for i := 0; i < n && i*2+1 < len(data); i++ {
			// 保持 PNG 兼容路径与 native display 的 16bpp 大端 RGB565 解析一致。 // Logan.Li
			c := binary.BigEndian.Uint16(data[i*2 : i*2+2])
			r := uint8((c >> 11) & 0x1F)
			g := uint8((c >> 5) & 0x3F)
			b := uint8(c & 0x1F)
			off := i * 4
			pix[off] = (r << 3) | (r >> 2)
			pix[off+1] = (g << 2) | (g >> 4)
			pix[off+2] = (b << 3) | (b >> 2)
			pix[off+3] = 0xFF
		}
	default:
		stride := rect.BitsPerPixel
		n := rect.Width * rect.Height
		for i := 0; i < n; i++ {
			s := i * stride
			if s+2 >= len(data) {
				break
			}
			off := i * 4
			pix[off+2] = data[s]
			pix[off+1] = data[s+1]
			pix[off] = data[s+2]
			pix[off+3] = 0xFF
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return "", err
	}
	return "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes()), nil
}

func binaryLittleUint16(b []byte) uint16 {
	return uint16(b[0]) | uint16(b[1])<<8
}