// Генерация ClientHello по логике Telegram Desktop (mtproto_tls_socket.cpp, PrepareClientHelloRules).
// См. https://github.com/telegramdesktop/tdesktop/blob/dev/Telegram/SourceFiles/mtproto/details/mtproto_tls_socket.cpp

package faketls

import (
	"crypto/ed25519"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"math/big"
	"time"
)

const (
	kMaxGrease        = 8
	kClientHelloLimit = 2048
	kHelloDigestLen   = 32
	kLengthSize       = 2
	kElementsM        = 384
	kAddedM           = 32
)

// BuildTdesktopClientHello строит запись TLS 1.2 ClientHello как в tdesktop (Ed25519, permutation, ECH padding, HMAC).
func BuildTdesktopClientHello(secretKey []byte, domain []byte) (*ClientHello, error) {
	if len(secretKey) != 16 {
		return nil, fmt.Errorf("secret key must be 16 bytes")
	}
	greases := prepareGreases()
	p := newTdesktopPart(domain, greases)
	writeClientHelloRules(p)
	if p.err {
		return nil, fmt.Errorf("tdesktop client hello: build failed")
	}
	p.finalize(secretKey)
	if p.err {
		return nil, fmt.Errorf("tdesktop client hello: finalize failed")
	}
	digest := p.extractDigest()
	out := p.take()
	if len(out) == 0 {
		return nil, fmt.Errorf("tdesktop client hello: empty buffer")
	}
	return &ClientHello{
		Record:      out,
		RandomField: append([]byte(nil), digest...),
		SessionID:   p.sessionID(),
	}, nil
}

func prepareGreases() []byte {
	result := make([]byte, kMaxGrease)
	if _, err := rand.Read(result); err != nil {
		return bytesRepeat(0x0a, kMaxGrease)
	}
	for i := range result {
		result[i] = (result[i] & 0xf0) + 0x0a
	}
	for i := 0; i < kMaxGrease; i += 2 {
		if result[i] == result[i+1] {
			result[i+1] ^= 0x10
		}
	}
	return result
}

func bytesRepeat(b byte, n int) []byte {
	out := make([]byte, n)
	for i := range out {
		out[i] = b
	}
	return out
}

type tdesktopPart struct {
	domain         []byte
	greases        []byte
	buf            []byte
	digestPosition int
	err            bool
}

func newTdesktopPart(domain []byte, greases []byte) *tdesktopPart {
	return &tdesktopPart{
		domain:         domain,
		greases:        greases,
		digestPosition: -1,
	}
}

func (p *tdesktopPart) grow(n int) []byte {
	if p.err || n <= 0 || len(p.buf)+n > kClientHelloLimit {
		p.err = true
		return nil
	}
	off := len(p.buf)
	p.buf = append(p.buf, make([]byte, n)...)
	return p.buf[off : off+n]
}

func (p *tdesktopPart) writeString(b []byte) {
	st := p.grow(len(b))
	if st == nil {
		return
	}
	copy(st, b)
}

func (p *tdesktopPart) writeZeros(n int) {
	already := len(p.buf)
	st := p.grow(n)
	if st == nil {
		return
	}
	if n == kHelloDigestLen && p.digestPosition < 0 {
		p.digestPosition = already
	}
	clear(st)
}

func (p *tdesktopPart) writeGrease(seed int) {
	if seed < 0 || seed >= len(p.greases) {
		p.err = true
		return
	}
	st := p.grow(2)
	if st == nil {
		return
	}
	// В tdesktop: bytes::set_with_const(storage, _greases[seed]) — оба байта равны значению seed.
	st[0] = p.greases[seed]
	st[1] = p.greases[seed]
}

func (p *tdesktopPart) writeRandom(n int) {
	st := p.grow(n)
	if st == nil {
		return
	}
	if _, err := rand.Read(st); err != nil {
		p.err = true
	}
}

func (p *tdesktopPart) writeDomain() {
	st := p.grow(len(p.domain))
	if st == nil {
		return
	}
	copy(st, p.domain)
}

func (p *tdesktopPart) writeEd25519PublicKey() {
	pub, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		p.err = true
		return
	}
	st := p.grow(32)
	if st == nil {
		return
	}
	copy(st, pub)
}

func (p *tdesktopPart) writeScope(f func()) {
	st := p.grow(kLengthSize)
	if st == nil {
		return
	}
	already := len(p.buf)
	f()
	length := len(p.buf) - already
	if length > 65535 {
		p.err = true
		return
	}
	binary.BigEndian.PutUint16(p.buf[already-2:already], uint16(length))
}

func (p *tdesktopPart) newSubPart() *tdesktopPart {
	return newTdesktopPart(p.domain, p.greases)
}

func (p *tdesktopPart) writePermutation(elements []func(*tdesktopPart)) {
	var parts [][]byte
	for _, el := range elements {
		sub := p.newSubPart()
		el(sub)
		if sub.err {
			p.err = true
			return
		}
		parts = append(parts, sub.take())
	}
	shuffleByteSlices(parts)
	for _, b := range parts {
		if p.err {
			return
		}
		st := p.grow(len(b))
		if st == nil {
			return
		}
		copy(st, b)
	}
}

func shuffleByteSlices(s [][]byte) {
	n := len(s)
	for i := n - 1; i > 0; i-- {
		jBig, err := rand.Int(rand.Reader, big.NewInt(int64(i+1)))
		j := 0
		if err != nil {
			j = i
		} else {
			j = int(jBig.Int64())
		}
		s[i], s[j] = s[j], s[i]
	}
}

func (p *tdesktopPart) writeM() {
	kElements := kElementsM
	kAdded := kAddedM
	storage := p.grow(kElements*3 + kAdded)
	if storage == nil {
		return
	}
	random := make([]byte, kElements*8+kAdded)
	if _, err := rand.Read(random); err != nil {
		p.err = true
		return
	}
	chars := storage
	ints := make([]uint32, len(random)/4)
	for i := 0; i < len(ints); i++ {
		ints[i] = binary.LittleEndian.Uint32(random[i*4 : i*4+4])
	}
	ci := 0
	for i := 0; i < kElements; i++ {
		a := int(ints[i*2] % 3329)
		b := int(ints[i*2+1] % 3329)
		chars[ci] = byte(a & 255)
		chars[ci+1] = byte((a >> 8) + ((b & 15) << 4))
		chars[ci+2] = byte(b >> 4)
		ci += 3
	}
	tail := storage[kElements*3:]
	if _, err := rand.Read(tail); err != nil {
		p.err = true
	}
}

func (p *tdesktopPart) writeE() {
	lengths := []int{144, 176, 208, 240}
	var pick [1]byte
	if _, err := rand.Read(pick[:]); err != nil {
		p.err = true
		return
	}
	length := lengths[int(pick[0])%len(lengths)]
	p.writeRandom(length)
}

func (p *tdesktopPart) writePadding() {
	if p.err {
		return
	}
	cur := len(p.buf)
	if cur >= 513 {
		return
	}
	z := 513 - cur
	p.writeString([]byte{0x00, 0x15})
	p.writeScope(func() {
		p.writeZeros(z)
	})
}

func (p *tdesktopPart) finalize(key []byte) {
	if p.err || p.digestPosition < 0 {
		p.err = true
		return
	}
	mac := hmac.New(sha256.New, key)
	mac.Write(p.buf)
	sum := mac.Sum(nil)
	copy(p.buf[p.digestPosition:p.digestPosition+kHelloDigestLen], sum)
	injectTimestamp(p.buf[p.digestPosition : p.digestPosition+kHelloDigestLen])
}

func injectTimestamp(digest32 []byte) {
	if len(digest32) != 32 {
		return
	}
	st := digest32[28:32]
	u := binary.LittleEndian.Uint32(st)
	ts := uint32(time.Now().Unix())
	u ^= ts
	binary.LittleEndian.PutUint32(st, u)
}

func (p *tdesktopPart) extractDigest() []byte {
	if p.digestPosition < 0 {
		return nil
	}
	return append([]byte(nil), p.buf[p.digestPosition:p.digestPosition+kHelloDigestLen]...)
}

func (p *tdesktopPart) take() []byte {
	if p.err {
		return nil
	}
	return append([]byte(nil), p.buf...)
}

// sessionID: первый 32-байтный Random сразу после префикса 0x20 (session id в ClientHello).
func (p *tdesktopPart) sessionID() []byte {
	if p.err {
		return nil
	}
	return firstR32Payload(p.buf)
}

func firstR32Payload(buf []byte) []byte {
	pat := []byte{0x03, 0x03}
	i := bytesIndex(buf, pat)
	if i < 0 {
		return nil
	}
	j := i + len(pat) + 32 // после Z(32) digest
	if j+1+32 > len(buf) {
		return nil
	}
	if buf[j] != 0x20 {
		return nil
	}
	return append([]byte(nil), buf[j+1:j+1+32]...)
}

func bytesIndex(hay []byte, needle []byte) int {
outer:
	for i := 0; i+len(needle) <= len(hay); i++ {
		for k := range needle {
			if hay[i+k] != needle[k] {
				continue outer
			}
		}
		return i
	}
	return -1
}

var cipherList32 = []byte{
	0x13, 0x01, 0x13, 0x02, 0x13, 0x03, 0xc0, 0x2b, 0xc0, 0x2f, 0xc0, 0x2c, 0xc0, 0x30,
	0xcc, 0xa9, 0xcc, 0xa8, 0xc0, 0x13, 0xc0, 0x14, 0x00, 0x9c, 0x00, 0x9d, 0x00, 0x2f, 0x00, 0x35,
	0x01, 0x00,
}

func writeClientHelloRules(p *tdesktopPart) {
	p.writeString([]byte{0x16, 0x03, 0x01})
	p.writeScope(func() {
		p.writeString([]byte{0x01, 0x00})
		p.writeScope(func() {
			p.writeString([]byte{0x03, 0x03})
			p.writeZeros(32)
			p.writeString([]byte{0x20})
			p.writeRandom(32)
			p.writeString([]byte{0x00, 0x20})
			p.writeGrease(0)
			p.writeString(cipherList32)
			p.writeScope(func() {
				p.writeGrease(2)
				p.writeString([]byte{0x00, 0x00})
				p.writePermutation(tdesktopPermutationElements())
				p.writeGrease(3)
				p.writeString([]byte{0x00, 0x01, 0x00})
				p.writePadding()
			})
		})
	})
}

func tdesktopPermutationElements() []func(*tdesktopPart) {
	return []func(*tdesktopPart){
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x00})
			p.writeScope(func() {
				p.writeScope(func() {
					p.writeString([]byte{0x00})
					p.writeScope(func() {
						p.writeDomain()
					})
				})
			})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x0a, 0x00, 0x0c, 0x00, 0x0a})
			p.writeGrease(4)
			p.writeString([]byte{0x11, 0xec, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x0b, 0x00, 0x02, 0x01, 0x00})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{
				0x00, 0x0d, 0x00, 0x12, 0x00, 0x10, 0x04, 0x03, 0x08, 0x04, 0x04, 0x01, 0x05, 0x03,
				0x08, 0x05, 0x05, 0x01, 0x08, 0x06, 0x06, 0x01,
			})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{
				0x00, 0x10, 0x00, 0x0e, 0x00, 0x0c, 0x02, 0x68, 0x32, 0x08, 0x68, 0x74, 0x74, 0x70,
				0x2f, 0x31, 0x2e, 0x31,
			})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x12, 0x00, 0x00})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x17, 0x00, 0x00})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x1b, 0x00, 0x03, 0x02, 0x00, 0x02})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x23, 0x00, 0x00})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x2b, 0x00, 0x07, 0x06})
			p.writeGrease(6)
			p.writeString([]byte{0x03, 0x04, 0x03, 0x03})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x2d, 0x00, 0x02, 0x01, 0x01})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x00, 0x33, 0x04, 0xef, 0x04, 0xed})
			p.writeGrease(4)
			p.writeString([]byte{0x00, 0x01, 0x00, 0x11, 0xec, 0x04, 0xc0})
			p.writeM()
			p.writeEd25519PublicKey()
			p.writeString([]byte{0x00, 0x1d, 0x00, 0x20})
			p.writeEd25519PublicKey()
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0x44, 0xcd, 0x00, 0x05, 0x00, 0x03, 0x02, 0x68, 0x32})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0xfe, 0x0d})
			p.writeScope(func() {
				p.writeString([]byte{0x00, 0x00, 0x01, 0x00, 0x01})
				p.writeRandom(1)
				p.writeString([]byte{0x00, 0x20})
				p.writeRandom(32)
				p.writeScope(func() {
					p.writeE()
				})
			})
		},
		func(p *tdesktopPart) {
			p.writeString([]byte{0xff, 0x01, 0x00, 0x01, 0x00})
		},
	}
}