package rdp

import (
	"bytes"
	"fmt"

	"ezops/pkg/rdp/trace"
)

const (
	fastPathOutputCompressionUsed = 0xC0
	fastPathFragmentSingle        = 0x00
	fastPathFragmentLast          = 0x10
	fastPathFragmentFirst         = 0x20
	fastPathFragmentNext          = 0x30
)

// FastPathReceiver 解析服务端 Fast-Path 更新并产出位图块 // Logan.Li
type FastPathReceiver struct {
	mppc        *MppcDecompressor
	fragments   bytes.Buffer
	frameIDs    []uint32
	updateStats map[byte]int
}

func NewFastPathReceiver() *FastPathReceiver {
	return &FastPathReceiver{mppc: NewMppcDecompressor(), updateStats: make(map[byte]int)}
}

func (f *FastPathReceiver) Handle(secFlag byte, payload []byte) ([]BitmapRect, error) {
	_ = secFlag
	r := bytes.NewReader(payload)
	var rects []BitmapRect
	for r.Len() > 0 {
		updateHeader, err := readU8(r)
		if err != nil {
			return rects, err
		}
		updateCode := updateHeader & 0x0f
		fragmentation := updateHeader & 0x30
		compression := updateHeader & 0xC0
		f.updateStats[updateCode]++

		var compressionFlags byte
		if compression == fastPathOutputCompressionUsed {
			compressionFlags, err = readU8(r)
			if err != nil {
				return rects, err
			}
		}
		size, err := readLE16(r)
		if err != nil {
			return rects, err
		}
		if f.updateStats[updateCode] <= 8 || f.updateStats[updateCode]%100 == 0 {
			trace.Logf("fastpath update code=%s(0x%x) count=%d size=%d frag=0x%02x compression=0x%02x flags=0x%02x", fastPathUpdateName(updateCode), updateCode, f.updateStats[updateCode], size, fragmentation, compression, compressionFlags)
		}
		block := make([]byte, int(size))
		if _, err = r.Read(block); err != nil {
			return rects, err
		}

		if fragmentation != fastPathFragmentSingle {
			if fragmentation == fastPathFragmentFirst {
				f.fragments.Reset()
			}
			f.fragments.Write(block)
			if fragmentation != fastPathFragmentLast {
				continue
			}
			block = append([]byte(nil), f.fragments.Bytes()...)
		}
		if compressionFlags != 0 {
			decompressed, err := f.mppc.Decompress(compressionFlags, block)
			if err != nil {
				continue
			}
			block = decompressed
		}

		switch updateCode {
		case fastPathUpdateBitmap:
			bs, err := parseFastPathBitmap(block)
			if err != nil {
				return rects, err
			}
			rects = append(rects, bs...)
		case fastPathUpdateSurfaceCmds:
			bs, frameIDs, err := parseSurfaceCommands(block)
			if err != nil {
				return rects, err
			}
			f.frameIDs = append(f.frameIDs, frameIDs...)
			rects = append(rects, bs...)
		default:
			trace.Logf("fastpath skip update code=%s(0x%x) size=%d", fastPathUpdateName(updateCode), updateCode, len(block))
		}
	}
	return rects, nil
}

func fastPathUpdateName(code byte) string {
	switch code {
	case 0x0:
		return "orders"
	case fastPathUpdateBitmap:
		return "bitmap"
	case 0x2:
		return "palette"
	case 0x3:
		return "synchronize"
	case fastPathUpdateSurfaceCmds:
		return "surface"
	case 0x5:
		return "ptr_null"
	case 0x6:
		return "ptr_default"
	case 0x8:
		return "ptr_position"
	case 0x9:
		return "color"
	case 0x0a:
		return "ptr_cached"
	case 0x0b:
		return "pointer"
	case 0x0c:
		return "large_pointer"
	default:
		return "unknown"
	}
}

// TakeFrameIDs 取出需要确认的 Surface frame，避免服务端因队列满停止推帧 // Logan.Li
func (f *FastPathReceiver) TakeFrameIDs() []uint32 {
	if len(f.frameIDs) == 0 {
		return nil
	}
	out := append([]uint32(nil), f.frameIDs...)
	f.frameIDs = f.frameIDs[:0]
	return out
}

func (f *FastPathReceiver) HandleSlowPath(payload []byte) ([]BitmapRect, error) {
	ind, err := parseMCSDataIndication(payload)
	if err != nil {
		return nil, err
	}
	if len(ind.payload) < 6 {
		return nil, nil
	}
	return parseSlowPathBitmap(ind.payload[4:])
}

func parseSurfaceCommands(data []byte) ([]BitmapRect, []uint32, error) {
	r := bytes.NewReader(data)
	var out []BitmapRect
	var frameIDs []uint32
	for r.Len() >= 2 {
		cmdType, err := readLE16(r)
		if err != nil {
			return out, frameIDs, err
		}
		switch cmdType {
		case 0x0001, 0x0006:
			wb, err := readSurfaceBits(r)
			if err != nil {
				return out, frameIDs, err
			}
			if rect, ok := decodeWireBitmap(wb); ok {
				out = append(out, rect)
			}
		case 0x0004:
			frameAction, err := readLE16(r)
			if err != nil {
				return out, frameIDs, err
			}
			frameID, err := readLE32(r)
			if err != nil {
				return out, frameIDs, err
			}
			trace.Logf("surface frame marker action=%d frameID=%d", frameAction, frameID)
			if frameAction == 0x0001 {
				frameIDs = append(frameIDs, frameID)
			}
		default:
			return out, frameIDs, fmt.Errorf("未知 SurfaceCmd: 0x%04x", cmdType)
		}
	}
	return out, frameIDs, nil
}

func readSurfaceBits(r *bytes.Reader) (wireBitmap, error) {
	destLeft, err := readLE16(r)
	if err != nil {
		return wireBitmap{}, err
	}
	destTop, err := readLE16(r)
	if err != nil {
		return wireBitmap{}, err
	}
	destRight, err := readLE16(r)
	if err != nil {
		return wireBitmap{}, err
	}
	destBottom, err := readLE16(r)
	if err != nil {
		return wireBitmap{}, err
	}
	bpp, err := readU8(r)
	if err != nil {
		return wireBitmap{}, err
	}
	flags, err := readU8(r)
	if err != nil {
		return wireBitmap{}, err
	}
	if _, err = readU8(r); err != nil {
		return wireBitmap{}, err
	}
	codecID, err := readU8(r)
	if err != nil {
		return wireBitmap{}, err
	}
	width, err := readLE16(r)
	if err != nil {
		return wireBitmap{}, err
	}
	height, err := readLE16(r)
	if err != nil {
		return wireBitmap{}, err
	}
	dataLen, err := readLE32(r)
	if err != nil {
		return wireBitmap{}, err
	}
	if flags&0x01 != 0 {
		if dataLen < 24 {
			return wireBitmap{}, fmt.Errorf("SurfaceBits 扩展头长度非法: %d", dataLen)
		}
		if _, err = r.Seek(24, 1); err != nil {
			return wireBitmap{}, err
		}
		dataLen -= 24
	}
	data := make([]byte, int(dataLen))
	if _, err = r.Read(data); err != nil {
		return wireBitmap{}, err
	}
	if codecID != 0 {
		trace.Logf("surface skip unsupported codec=%d rect=%dx%d@%d,%d data=%d flags=0x%02x", codecID, width, height, destLeft, destTop, dataLen, flags)
		return wireBitmap{}, nil
	}
	trace.Logf("surface bits codec=%d rect=%dx%d@%d,%d data=%d flags=0x%02x", codecID, width, height, destLeft, destTop, dataLen, flags)
	return wireBitmap{
		DestLeft: destLeft, DestTop: destTop, DestRight: destRight, DestBottom: destBottom,
		Width: width, Height: height, BitsPerPixel: uint16(bpp), Flags: uint16(flags) | bitmapNoProcessing,
		BitmapDataStream: data,
	}, nil
}

func readU8(r *bytes.Reader) (byte, error) {
	b, err := r.ReadByte()
	return b, err
}