package core

import (
	"errors"
	"fmt"
)

// MppcDecompressor maintains per-connection state for RDP bulk data
// decompression (MS-RDPBCGR §3.1.8.4, RFC 2118 MPPC token grammar).  A single
// instance is shared between the fast-path and slow-path receivers of one RDP
// connection.  The token grammar follows the MS-RDPBCGR pseudo-code as
// implemented in FreeRDP's libfreerdp/codec/mppc.c (rdesktop's variant uses a
// different literal/length encoding that does not interop with Win10 RDP5
// streams).
//
// Token grammar (bit-level, MSB first):
//
//	literal 0x00-0x7F : "0"  + 7 bits (8-bit token)
//	literal 0x80-0xFF : "10" + 7 bits (9-bit token)
//	copy tuple        : "11" + offset prefix + length prefix
//
// Offset prefix (RDP5, compressionType nibble 0x1, 64K dictionary):
//
//	"11111" + 6 bits  → offset 0-63
//	"11110" + 8 bits  → offset 64-319
//	"1110"  + 11 bits → offset 320-2367
//	"110"   + 16 bits → offset 2368-67903
//
// (RDP4 / 8K dictionary: "1111"+6, "1110"+8, "110"+13.)
//
// Length prefix ("0" → 3; otherwise n 1-bits, a 0 terminator, then n+1 value
// bits with the (1<<m) high bit implied, m = n+1; RDP5 allows m up to 15).
//
// History: uncompressed segments do NOT touch the shared history; FLUSH
// (0x80) zeroes it, RESET/AT_FRONT (0x40) rewinds the write cursor to 0.
type MppcDecompressor struct {
	history [mppcHistorySize]byte
	offset  int
}

const mppcHistorySize = 65536

// mppc flag bits（fast-path compressionFlags 字节，MS-RDPBCGR 2.2.9.1.1.4.1
// + FreeRDP bulk.c）：
//
//	低半字节 0x01      – 压缩类型选择：1 = RDP5（64K 字典）
//	0x20 PACKET_COMPRESSED – 块为 MPPC 压缩数据
//	0x40 PACKET_AT_FRONT    – 解压数据写到历史前端（游标归零）
//	0x80 PACKET_FLUSHED     – 清零历史，游标归零
//
// 注意：FreeRDP mppc.c 不处理 0x01 位（它只是类型选择），本解码器固定
// 使用 64K 字典，因此该位可忽略。
const (
	mppcType64K    = 0x01
	mppcCompressed = 0x20
	mppcAtFront    = 0x40
	mppcFlushed    = 0x80
)

func NewMppcDecompressor() *MppcDecompressor {
	return &MppcDecompressor{}
}

// Decompress processes one MPPC segment.
//
// flags is the compressionFlags byte (fast-path) or the slow-path bulk flags;
// it carries the PACKET_* bits.  Segments without PACKET_COMPRESSED are
// returned unchanged and do not affect the history.
func (d *MppcDecompressor) Decompress(flags byte, data []byte) ([]byte, error) {
	if flags&mppcFlushed != 0 {
		d.history = [mppcHistorySize]byte{}
		d.offset = 0
	}
	if flags&mppcAtFront != 0 {
		d.offset = 0
	}

	if flags&mppcCompressed == 0 {
		out := make([]byte, len(data))
		copy(out, data)
		return out, nil
	}

	// RDP5 64K 字典（协商固定，见常量注释）。
	const mask = 65535
	start := d.offset
	br := newMppcBitReader(data)

	// abort 与 FreeRDP mppc.c 失败路径语义一致：HistoryPtr 只在成功返回时
	// 提交，解码中途失败必须回滚写游标，否则后续分段的 match 引用的
	// 历史从此错位，整条流持续产出垃圾。
	abort := func(msg string) ([]byte, error) {
		d.offset = start
		return nil, errors.New("mppc: " + msg)
	}

	for br.bitsLeft >= 8 {
		if br.readBit() == 0 {
			// Literal 0x00-0x7F ("0" + 7 bits).
			if d.offset >= mppcHistorySize {
				return abort("history full")
			}
			d.history[d.offset] = byte(br.readBits(7))
			d.offset++
			continue
		}
		if br.readBit() == 0 {
			// Literal 0x80-0xFF ("10" + 7 bits, 9-bit token).
			if d.offset >= mppcHistorySize {
				return abort("history full")
			}
			d.history[d.offset] = byte(0x80 | br.readBits(7))
			d.offset++
			continue
		}

		// Copy tuple: decode CopyOffset (distance back from the write
		// cursor, masked into the dictionary)。前缀位必须逐位惰性读取，
		// 不能在 switch 初始化里预先消耗。
		var copyOffset int
		if br.readBit() == 0 {
			copyOffset = br.readBits(16) + 2368
		} else if br.readBit() == 0 {
			copyOffset = br.readBits(11) + 320
		} else if br.readBit() == 0 {
			copyOffset = br.readBits(8) + 64
		} else {
			copyOffset = br.readBits(6)
		}

		// Decode LengthOfMatch: n leading 1-bits + 0 terminator, then n+1
		// value bits; length = (1<<(n+1)) | bits.  "0" alone means 3.
		n := 0
		for br.readBit() == 1 {
			n++
			const maxBits = 15
			if n > maxBits-1 {
				return abort("length code overflow")
			}
		}
		var copyLength int
		if n == 0 {
			copyLength = 3
		} else {
			m := n + 1
			copyLength = (1 << uint(m)) | br.readBits(m)
		}

		if d.offset+copyLength > mppcHistorySize {
			return abort(fmt.Sprintf("copy overflows history (%d+%d)", d.offset, copyLength))
		}
		src := (d.offset - copyOffset) & mask
		for i := 0; i < copyLength; i++ {
			d.history[d.offset] = d.history[src]
			d.offset++
			src = (src + 1) & mask
		}
	}

	out := make([]byte, d.offset-start)
	copy(out, d.history[start:d.offset])
	return out, nil
}

// mppcBitReader reads bits MSB-first from a byte slice.  Reading past the
// end yields zero bits (matching FreeRDP's padded bit stream behaviour).
type mppcBitReader struct {
	data     []byte
	byteIdx  int
	mask     byte // bit mask within current byte; starts at 0x80
	bitsLeft int  // total bits remaining
}

func newMppcBitReader(data []byte) *mppcBitReader {
	return &mppcBitReader{
		data:     data,
		mask:     0x80,
		bitsLeft: len(data) * 8,
	}
}

func (r *mppcBitReader) readBit() int {
	if r.bitsLeft <= 0 {
		return 0
	}
	r.bitsLeft--
	var bit int
	if r.data[r.byteIdx]&r.mask != 0 {
		bit = 1
	}
	r.mask >>= 1
	if r.mask == 0 {
		r.mask = 0x80
		r.byteIdx++
	}
	return bit
}

func (r *mppcBitReader) readBits(n int) int {
	result := 0
	for i := 0; i < n; i++ {
		result = (result << 1) | r.readBit()
	}
	return result
}