package mtproxy

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	crand "crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"io"
)

// ObfuscateTagRandomizedIntermediate matches Telethon RandomizedIntermediatePacketCodec.
var ObfuscateTagRandomizedIntermediate = []byte{0xdd, 0xdd, 0xdd, 0xdd}

// randSource is injectable for tests.
var randSource io.Reader = crand.Reader

// InitHeader returns the 64-byte first payload (MTProxyIO.init_header) and CTR streams.
func InitHeader(secret16 []byte, dcID int16) (header64 []byte, enc cipher.Stream, dec cipher.Stream, err error) {
	if len(secret16) != 16 {
		return nil, nil, nil, fmt.Errorf("mtproxy secret must be 16 bytes")
	}
	sec := secret16

	keywords := [][]byte{
		[]byte("PVrG"),
		[]byte("GET "),
		[]byte("POST"),
		{0xee, 0xee, 0xee, 0xee},
	}
	var rnd [64]byte
	for {
		if _, err := io.ReadFull(randSource, rnd[:]); err != nil {
			return nil, nil, nil, err
		}
		if rnd[0] == 0xef {
			continue
		}
		bad := false
		for _, k := range keywords {
			if len(k) <= 4 && bytes.Equal(rnd[:len(k)], k) {
				bad = true
				break
			}
		}
		if bad {
			continue
		}
		if rnd[4] == 0 && rnd[5] == 0 && rnd[6] == 0 && rnd[7] == 0 {
			continue
		}
		break
	}

	random := rnd[:]
	rev := make([]byte, 48)
	for i := 0; i < 48; i++ {
		rev[i] = random[55-i]
	}

	encKey := sha256Sum(append(append([]byte{}, random[8:40]...), sec...))
	encIV := random[40:56]
	decKey := sha256Sum(append(append([]byte{}, rev[:32]...), sec...))
	decIV := rev[32:48]

	enc, err = newCTR(encKey, encIV)
	if err != nil {
		return nil, nil, nil, err
	}
	dec, err = newCTR(decKey, decIV)
	if err != nil {
		return nil, nil, nil, err
	}

	buf := make([]byte, 64)
	copy(buf, random)
	copy(buf[56:60], ObfuscateTagRandomizedIntermediate)
	dcBytes := make([]byte, 2)
	binary.LittleEndian.PutUint16(dcBytes, uint16(dcID))
	copy(buf[60:62], dcBytes)

	encChunk := make([]byte, 8)
	copy(encChunk, buf[56:64])
	enc.XORKeyStream(encChunk, encChunk)
	copy(buf[56:64], encChunk)

	return buf, enc, dec, nil
}

func sha256Sum(b []byte) []byte {
	h := sha256.Sum256(b)
	return h[:]
}

func newCTR(key, iv []byte) (cipher.Stream, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewCTR(block, iv), nil
}