package faketls

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"io"
)

// Маркеры и разбор длин — как в mtproto_tls_socket.cpp (checkHelloParts12 / checkHelloParts34),
// а не фиксированные 127+6+3+2 байта telethon + «http»-хвост (на части хостов это давало ложную длину хвоста).
var (
	kServerHelloPart1 = []byte{0x16, 0x03, 0x03}
	kServerHelloPart3 = []byte{0x14, 0x03, 0x03, 0x00, 0x01, 0x01, 0x17, 0x03, 0x03}
)

const kTLSLengthSize = 2

// ReadServerHello читает цепочку fake-TLS ServerHello до конца части, по которой считается digest (как в tdesktop).
func ReadServerHello(r io.Reader) ([]byte, error) {
	parts1Size := len(kServerHelloPart1) + kTLSLengthSize
	header := make([]byte, parts1Size)
	if _, err := io.ReadFull(r, header); err != nil {
		return nil, fmt.Errorf("read tls record header: %w", err)
	}
	part1Off := parts1Size - kTLSLengthSize - len(kServerHelloPart1)
	if !bytes.Equal(header[part1Off:part1Off+len(kServerHelloPart1)], kServerHelloPart1) {
		return nil, fmt.Errorf("server hello: bad tls prefix")
	}
	part2Size := int(binary.BigEndian.Uint16(header[parts1Size-kTLSLengthSize:]))
	if part2Size <= 0 || part2Size > 1<<20 {
		return nil, fmt.Errorf("invalid tls fragment length %d", part2Size)
	}
	midLen := part2Size + len(kServerHelloPart3) + kTLSLengthSize
	mid := make([]byte, midLen)
	if _, err := io.ReadFull(r, mid); err != nil {
		return nil, fmt.Errorf("read server hello body: %w", err)
	}
	part3Off := part2Size
	if !bytes.Equal(mid[part3Off:part3Off+len(kServerHelloPart3)], kServerHelloPart3) {
		return nil, fmt.Errorf("server hello: unexpected middle tls markers")
	}
	part4Size := int(binary.BigEndian.Uint16(mid[len(mid)-kTLSLengthSize:]))
	if part4Size < 0 || part4Size > 1<<20 {
		return nil, fmt.Errorf("invalid final fragment length %d", part4Size)
	}
	tail := make([]byte, part4Size)
	if _, err := io.ReadFull(r, tail); err != nil {
		return nil, fmt.Errorf("read server hello final fragment: %w", err)
	}
	out := append(append(append([]byte{}, header...), mid...), tail...)
	return out, nil
}

// VerifyServerHelloFull validates the proxy response (telethon-faketls rules).
func VerifyServerHelloFull(full []byte, secretKey []byte, clientRandom []byte, sessionID []byte) error {
	if len(full) < 127+6+3+2 {
		return fmt.Errorf("server hello too short")
	}
	if !bytes.HasPrefix(full, []byte{0x16, 0x03, 0x03}) {
		return fmt.Errorf("server hello: bad tls prefix")
	}
	marker := []byte{0x14, 0x03, 0x03, 0x00, 0x01, 0x01, 0x17, 0x03, 0x03}
	if len(full) < 127+len(marker) || !bytes.Equal(full[127:127+len(marker)], marker) {
		return fmt.Errorf("server hello: unexpected middle tls markers")
	}
	sidStart := 11 + 32 + 1
	if len(full) < sidStart+32 || !bytes.Equal(full[sidStart:sidStart+32], sessionID) {
		return fmt.Errorf("server hello: session id mismatch")
	}
	patched := make([]byte, len(full))
	copy(patched, full)
	clear(patched[11:43])
	mac := hmac.New(sha256.New, secretKey)
	mac.Write(clientRandom)
	mac.Write(patched)
	want := mac.Sum(nil)
	if !bytes.Equal(full[11:43], want) {
		return fmt.Errorf("server hello: digest mismatch")
	}
	return nil
}

const (
	kServerHelloDigestPosition = 11
)

// VerifyServerHelloTdesktop проверяет ответ прокси по правилам mtproto_tls_socket.cpp (checkHelloDigest):
// HMAC-SHA256(key, client_digest||server_bytes) при обнулённом digest сервера в буфере.
func VerifyServerHelloTdesktop(serverResponse []byte, secretKey []byte, clientDigest []byte) error {
	if len(secretKey) != 16 {
		return fmt.Errorf("secret key must be 16 bytes")
	}
	if len(clientDigest) != 32 {
		return fmt.Errorf("client digest must be 32 bytes")
	}
	if len(serverResponse) < kServerHelloDigestPosition+32 {
		return fmt.Errorf("server hello too short for digest")
	}
	fulldata := make([]byte, 0, len(clientDigest)+len(serverResponse))
	fulldata = append(fulldata, clientDigest...)
	fulldata = append(fulldata, serverResponse...)
	digestOff := len(clientDigest) + kServerHelloDigestPosition
	digestCopy := append([]byte(nil), fulldata[digestOff:digestOff+32]...)
	clear(fulldata[digestOff : digestOff+32])
	mac := hmac.New(sha256.New, secretKey)
	mac.Write(fulldata)
	want := mac.Sum(nil)
	if !bytes.Equal(digestCopy, want) {
		return fmt.Errorf("server hello: digest mismatch (tdesktop)")
	}
	return nil
}