// Package tgquick performs a minimal MTProto step (req_pq → ответ DC) through an established MTProxy tunnel,
// matching Telethon's RandomizedIntermediate + MTProxyIO framing.
package tgquick

import (
	"bufio"
	"bytes"
	"context"
	"crypto/cipher"
	crand "crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"time"

	"mtproxy_checker/internal/faketls"
)

const (
	tlReqPQ = 0xd712e4be
	tlResPQ = 0x05162463
	maxRI   = 1 << 20
)

// ErrNoResPQ means strict deep check did not see resPQ#05162463 from DC.
var ErrNoResPQ = errors.New("no resPQ from telegram through proxy (tunnel may be broken)")

// ErrNoDCReply means relaxed deep check got no parsable unencrypted MTProto reply from DC before deadline.
var ErrNoDCReply = errors.New("no reply from telegram DC through proxy (timeout)")

// ErrPeerClosed is returned when the remote side closes TCP during the MTProto probe.
var ErrPeerClosed = errors.New("connection closed by peer during mtproto probe")

// ErrCleartextHTTP means the peer sent a plaintext HTTP line instead of continuing fake-TLS (often a front nginx 400/502).
var ErrCleartextHTTP = errors.New("cleartext http on socket (not fake-tls mtproxy payload)")

func newMessageID() int64 {
	t := time.Now().UnixNano()
	return (t / 4) &^ 3
}

func buildReqPQ() ([]byte, error) {
	nonce := make([]byte, 16)
	if _, err := io.ReadFull(crand.Reader, nonce); err != nil {
		return nil, err
	}
	body := make([]byte, 4+16)
	binary.LittleEndian.PutUint32(body, tlReqPQ)
	copy(body[4:], nonce)
	return wrapUnencrypted(body), nil
}

func wrapUnencrypted(msgData []byte) []byte {
	out := make([]byte, 8+8+4+len(msgData))
	// auth_key_id = 0
	binary.LittleEndian.PutUint64(out[8:], uint64(newMessageID()))
	binary.LittleEndian.PutUint32(out[16:], uint32(len(msgData)))
	copy(out[20:], msgData)
	return out
}

// randomizedIntermediateEncode packs payload with 0–3 random padding bytes (Telethon RandomizedIntermediatePacketCodec).
func randomizedIntermediateEncode(payload []byte) ([]byte, error) {
	var rnd [1]byte
	if _, err := crand.Read(rnd[:]); err != nil {
		return nil, err
	}
	pad := int(rnd[0] % 4)
	padding := make([]byte, pad)
	if _, err := crand.Read(padding); err != nil {
		return nil, err
	}
	body := append(append([]byte{}, payload...), padding...)
	out := make([]byte, 4+len(body))
	binary.LittleEndian.PutUint32(out, uint32(len(body)))
	copy(out[4:], body)
	return out, nil
}

func popRI(acc *[]byte) ([]byte, bool) {
	a := *acc
	if len(a) < 4 {
		return nil, false
	}
	L := int(binary.LittleEndian.Uint32(a[0:4]))
	if L < 0 || L > maxRI || len(a) < 4+L {
		return nil, false
	}
	frame := a[4 : 4+L]
	*acc = append([]byte(nil), a[4+L:]...)
	pad := len(frame) % 4
	if pad > 0 {
		frame = frame[:len(frame)-pad]
	}
	return frame, true
}

func isResPQ(mtInner []byte) bool {
	if len(mtInner) < 20 {
		return false
	}
	if binary.LittleEndian.Uint64(mtInner[0:8]) != 0 {
		return false
	}
	ml := int(binary.LittleEndian.Uint32(mtInner[16:20]))
	if ml < 4 || ml > maxRI || 20+ml > len(mtInner) {
		return false
	}
	body := mtInner[20 : 20+ml]
	return binary.LittleEndian.Uint32(body[0:4]) == tlResPQ
}

// isUnencryptedDCEnvelope reports a valid MTProto unencrypted container (auth_key_id 0) with non-empty body.
// After req_pq the DC normally sends resPQ, but any well-formed unencrypted reply proves the tunnel carried DC traffic.
func isUnencryptedDCEnvelope(mtInner []byte) bool {
	if len(mtInner) < 24 {
		return false
	}
	if binary.LittleEndian.Uint64(mtInner[0:8]) != 0 {
		return false
	}
	ml := int(binary.LittleEndian.Uint32(mtInner[16:20]))
	if ml < 4 || ml > maxRI || 20+ml > len(mtInner) {
		return false
	}
	return true
}

// DrainPostInitEE reads inbound fake-TLS records after MTProxy init and before req_pq (Telethon TcpMTProxy waits for data ~2s).
// Consumes 0x17 payloads with dec so the CTR stream stays aligned with the server; discards MTProto frames until idle.
func DrainPostInitEE(ctx context.Context, br *bufio.Reader, conn net.Conn, dec cipher.Stream, maxWait time.Duration) error {
	defer func() { _ = conn.SetReadDeadline(time.Time{}) }()
	end := time.Now().Add(maxWait)
	if d, ok := ctx.Deadline(); ok && d.Before(end) {
		end = d
	}
	var acc []byte
	const maxDrain = 256 << 10
	drained := 0
	for drained < maxDrain && time.Now().Before(end) {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		d := 500 * time.Millisecond
		if rem := time.Until(end); rem < d {
			d = rem
		}
		if cd, ok := ctx.Deadline(); ok {
			if rem := time.Until(cd); rem < d {
				d = rem
			}
		}
		if d <= 0 {
			return nil
		}
		_ = conn.SetReadDeadline(time.Now().Add(d))
		if err := sniffCleartextHTTP(br); err != nil {
			if ne, ok := err.(net.Error); ok && ne.Timeout() {
				continue
			}
			return err
		}
		chunk, err := readNextTLS17Payload(br)
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Timeout() {
				return nil
			}
			if errors.Is(err, ErrCleartextHTTP) {
				return err
			}
			if errors.Is(err, io.EOF) {
				return nil
			}
			return fmt.Errorf("post-init drain: %w", err)
		}
		dec.XORKeyStream(chunk, chunk)
		drained += len(chunk)
		acc = append(acc, chunk...)
		for {
			frame, ok := popRI(&acc)
			if !ok {
				break
			}
			_ = frame
		}
	}
	return nil
}

// VerifyResPQ sends req_pq and waits for a DC reply on the MTProxy byte stream (decrypt with dec, encrypt with enc).
// If strictResPQ is true, the first matching RI frame must be resPQ; otherwise any valid unencrypted MTProto message is enough.
// For ee (fake-TLS), pass the same bufio.Reader used after init (and DrainPostInitEE); for dd pass tlsBR=nil.
func VerifyResPQ(ctx context.Context, conn net.Conn, enc, dec cipher.Stream, tlsBR *bufio.Reader, strictResPQ bool) error {
	plain, err := buildReqPQ()
	if err != nil {
		return fmt.Errorf("build req_pq: %w", err)
	}
	ri, err := randomizedIntermediateEncode(plain)
	if err != nil {
		return fmt.Errorf("randomized intermediate: %w", err)
	}
	wire := append([]byte(nil), ri...)
	enc.XORKeyStream(wire, wire)
	if tlsBR != nil {
		if err := faketls.WriteTLSApplicationData(conn, wire); err != nil {
			return fmt.Errorf("write tls app (req_pq): %w", err)
		}
	} else {
		if _, err := conn.Write(wire); err != nil {
			return fmt.Errorf("write req_pq: %w", err)
		}
	}

	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(60 * time.Second)
	}

	var acc []byte

	const maxAcc = 256 << 10
	// Один read ждёт ответ DC; 3s мало при нагрузке/гео — держим до 30s, но не дольше ctx.
	const perReadSlack = 30 * time.Second
	for time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		chunkDeadline := time.Now().Add(perReadSlack)
		if chunkDeadline.After(deadline) {
			chunkDeadline = deadline
		}
		if !chunkDeadline.After(time.Now().Add(50 * time.Millisecond)) {
			return fmt.Errorf("%w: not enough time left waiting for resPQ", context.DeadlineExceeded)
		}
		_ = conn.SetReadDeadline(chunkDeadline)
		var chunk []byte
		var rerr error
		if tlsBR != nil {
			if err := sniffCleartextHTTP(tlsBR); err != nil {
				if ne, ok := err.(net.Error); ok && ne.Timeout() {
					if time.Now().Before(deadline) {
						continue
					}
					break
				}
				return err
			}
			chunk, rerr = readNextTLS17Payload(tlsBR)
		} else {
			chunk, rerr = readAtMost(conn, 65536)
		}
		if rerr != nil {
			if ne, ok := rerr.(net.Error); ok && ne.Timeout() {
				if time.Now().Before(deadline) {
					continue
				}
				break
			}
			if errors.Is(rerr, io.EOF) {
				return fmt.Errorf("%w", ErrPeerClosed)
			}
			if errors.Is(rerr, ErrCleartextHTTP) {
				return rerr
			}
			return fmt.Errorf("read mtproxy stream: %w", rerr)
		}
		if len(chunk) == 0 {
			continue
		}
		dec.XORKeyStream(chunk, chunk)
		acc = append(acc, chunk...)
		if len(acc) > maxAcc {
			return fmt.Errorf("mtproxy read buffer overflow")
		}
		for {
			frame, ok := popRI(&acc)
			if !ok {
				break
			}
			if strictResPQ {
				if isResPQ(frame) {
					_ = conn.SetReadDeadline(time.Time{})
					return nil
				}
				continue
			}
			if isUnencryptedDCEnvelope(frame) {
				_ = conn.SetReadDeadline(time.Time{})
				return nil
			}
		}
	}
	_ = conn.SetReadDeadline(time.Time{})
	if strictResPQ {
		return ErrNoResPQ
	}
	return ErrNoDCReply
}

// readNextTLS17Payload reads TLS 1.2-style records from the wire (plaintext record headers).
// Handshake (0x16) and ChangeCipherSpec (0x15) payloads are not MTProxy-CTR ciphertext — skip them.
// Only application data (0x17) inner bytes are passed through the MTProxy decrypt stream (caller XORs).
// Some proxies send extra handshake records after the fake ServerHello; clients skip them before MTProto.
func sniffCleartextHTTP(br *bufio.Reader) error {
	p, err := br.Peek(8)
	if len(p) >= 4 && bytes.HasPrefix(p, []byte("HTTP")) {
		line := string(p)
		if idx := bytes.IndexByte(p, '\n'); idx >= 0 {
			line = string(p[:idx])
		}
		return fmt.Errorf("%w: %s", ErrCleartextHTTP, line)
	}
	if err != nil && !errors.Is(err, bufio.ErrBufferFull) && err != io.EOF {
		return err
	}
	return nil
}

func readNextTLS17Payload(br *bufio.Reader) ([]byte, error) {
	const maxSkips = 128
	for range maxSkips {
		if err := sniffCleartextHTTP(br); err != nil {
			return nil, err
		}
		h := make([]byte, 5)
		if _, err := io.ReadFull(br, h); err != nil {
			return nil, err
		}
		ver := binary.BigEndian.Uint16(h[1:3])
		if ver != 0x0303 && ver != 0x0301 {
			if h[0] == 'H' && h[1] == 'T' {
				return nil, fmt.Errorf("%w: misaligned or non-tls data (type 0x%02x version 0x%04x)", ErrCleartextHTTP, h[0], ver)
			}
			return nil, fmt.Errorf("unexpected tls record version 0x%04x (type 0x%02x)", ver, h[0])
		}
		n := int(binary.BigEndian.Uint16(h[3:5]))
		if n <= 0 || n > 1<<20 {
			return nil, fmt.Errorf("invalid tls record length %d", n)
		}
		payload := make([]byte, n)
		if _, err := io.ReadFull(br, payload); err != nil {
			return nil, err
		}
		switch h[0] {
		case 0x17: // application data — inner bytes are MTProxy-obfuscated
			return payload, nil
		case 0x16, 0x15: // more handshake / CCS after ServerHello — plaintext, do not advance CTR interpretation here
			continue
		case 0x14: // alert
			return nil, fmt.Errorf("tls alert record from proxy (len=%d)", n)
		default:
			return nil, fmt.Errorf("unexpected tls record type 0x%02x (len=%d); 0x48 often means cleartext HTTP if the stream is misaligned", h[0], n)
		}
	}
	return nil, fmt.Errorf("too many non-application tls records before 0x17")
}

func readAtMost(conn net.Conn, max int) ([]byte, error) {
	buf := make([]byte, max)
	n, err := conn.Read(buf)
	if n > 0 {
		if err == io.EOF {
			err = nil
		}
		return buf[:n], err
	}
	return nil, err
}