package faketls

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

// Extensions inside ClientHello must total exactly 403 bytes (TelethonFakeTLS / MTProxy expectation).
const extensionsTotal = 403

// Fixed TLS 1.2 cipher list from telethon-faketls.
var cipherSuites = []byte{
	0xfa, 0xfa, 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,
}

var extBeforeKeyShare = []byte{
	0x00, 0x17, 0x00, 0x00,
	0xff, 0x01, 0x00, 0x01, 0x00,
	0x00, 0x0a, 0x00, 0x0a, 0x00, 0x08, 0xba, 0xba, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18,
	0x00, 0x0b, 0x00, 0x02, 0x01, 0x00,
	0x00, 0x23, 0x00, 0x00,
	0x00, 0x10, 0x00, 0x0e, 0x00, 0x0c, 0x02, 0x68, 0x32, 0x08, 0x68, 0x74, 0x74, 0x70, 0x2f, 0x31, 0x2e, 0x31,
	0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x0d, 0x00, 0x12, 0x00, 0x10, 0x04, 0x03, 0x08, 0x04, 0x04, 0x01, 0x05, 0x03, 0x08, 0x05, 0x05, 0x01, 0x08, 0x06, 0x06, 0x01,
	0x00, 0x12, 0x00, 0x00,
}

var keyShareHeader = []byte{0x00, 0x33, 0x00, 0x2b, 0x00, 0x29, 0xba, 0xba, 0x00, 0x01, 0x00, 0x00, 0x1d, 0x00, 0x20}

var extAfterKeyShare = []byte{
	0x00, 0x2d, 0x00, 0x02, 0x01, 0x01,
	0x00, 0x2b, 0x00, 0x0b, 0x0a, 0x9a, 0x9a, 0x03, 0x04, 0x03, 0x03, 0x03, 0x02, 0x03, 0x01,
	0x00, 0x1b, 0x00, 0x03, 0x02, 0x00, 0x02,
	0x1a, 0x1a, 0x00, 0x01, 0x00,
}

func sniExtension(domain []byte) []byte {
	entry := make([]byte, 0, 3+len(domain))
	entry = append(entry, 0x00)
	entry = append(entry, byte(len(domain)>>8), byte(len(domain)))
	entry = append(entry, domain...)
	list := make([]byte, 0, 2+len(entry))
	list = append(list, byte(len(entry)>>8), byte(len(entry)))
	list = append(list, entry...)
	out := make([]byte, 0, 4+len(list))
	out = append(out, 0x00, 0x00)
	out = append(out, byte(len(list)>>8), byte(len(list)))
	out = append(out, list...)
	return out
}

func buildExtensions(domain []byte, pubKey []byte) ([]byte, error) {
	if len(pubKey) != 32 {
		return nil, fmt.Errorf("key share pubkey: need 32 bytes")
	}
	var exts bytes.Buffer
	exts.Write([]byte{0x4a, 0x4a, 0x00, 0x00})
	exts.Write(sniExtension(domain))
	exts.Write(extBeforeKeyShare)
	exts.Write(keyShareHeader)
	exts.Write(pubKey)
	exts.Write(extAfterKeyShare)

	core := exts.Len()
	// padding extension: 00 15 + u16(len) + zeros; total extensions must be 403
	padHdr := 4
	padData := extensionsTotal - core - padHdr
	if padData < 0 {
		return nil, fmt.Errorf("mtproxy fake-tls: domain too long for fixed 403-byte extension block (core=%d)", core)
	}
	exts.Write([]byte{0x00, 0x15})
	exts.Write([]byte{byte(padData >> 8), byte(padData)})
	for i := 0; i < padData; i++ {
		exts.WriteByte(0x00)
	}
	if exts.Len() != extensionsTotal {
		return nil, fmt.Errorf("internal: extensions len %d != %d", exts.Len(), extensionsTotal)
	}
	return exts.Bytes(), nil
}

// ClientHello is the initial raw TLS record and the 32-byte ClientHello.random field (post-HMAC).
type ClientHello struct {
	Record      []byte
	RandomField []byte
	SessionID   []byte
}

// BuildClientHello constructs the first wire record (517 bytes) compatible with telethon-faketls.
func BuildClientHello(secretKey []byte, domain []byte) (*ClientHello, error) {
	if len(secretKey) != 16 {
		return nil, fmt.Errorf("secret key must be 16 bytes")
	}
	pub, err := GenX25519PublicKey()
	if err != nil {
		return nil, err
	}
	sessionID, err := Rand32()
	if err != nil {
		return nil, err
	}
	exts, err := buildExtensions(domain, pub)
	if err != nil {
		return nil, err
	}

	// ClientHello body = 508 bytes
	var body bytes.Buffer
	body.Write([]byte{0x03, 0x03})
	randomZero := make([]byte, 32)
	body.Write(randomZero)
	body.WriteByte(32)
	body.Write(sessionID)
	body.Write([]byte{0x00, 0x20})
	body.Write(cipherSuites)
	body.WriteByte(0x01)
	body.WriteByte(0x00)
	body.Write([]byte{byte(extensionsTotal >> 8), byte(extensionsTotal & 0xff)})
	body.Write(exts)
	if body.Len() != 508 {
		return nil, fmt.Errorf("internal: client hello body %d != 508", body.Len())
	}

	handshake := make([]byte, 0, 4+508)
	handshake = append(handshake, 0x01)
	handshake = append(handshake, byte(508>>16), byte((508>>8)&0xff), byte(508&0xff))
	handshake = append(handshake, body.Bytes()...)

	record := make([]byte, 0, 5+len(handshake))
	record = append(record, 0x16, 0x03, 0x01)
	record = binary.BigEndian.AppendUint16(record, uint16(len(handshake)))
	record = append(record, handshake...)

	randomField := computeRandomField(secretKey, record)
	copy(record[11:43], randomField)

	return &ClientHello{Record: record, RandomField: randomField, SessionID: sessionID}, nil
}

func computeRandomField(secretKey []byte, record []byte) []byte {
	msg := make([]byte, 0, len(record))
	msg = append(msg, record[:11]...)
	msg = append(msg, make([]byte, 32)...)
	msg = append(msg, record[43:]...)
	mac := hmac.New(sha256.New, secretKey)
	mac.Write(msg)
	digest := mac.Sum(nil)
	ts := uint32(time.Now().Unix())
	tsb := []byte{byte(ts), byte(ts >> 8), byte(ts >> 16), byte(ts >> 24)}
	out := make([]byte, 32)
	copy(out[:28], digest[:28])
	for i := 0; i < 4; i++ {
		out[28+i] = digest[28+i] ^ tsb[i]
	}
	return out
}