package mtprotoedge

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/go-faster/errors"
	"go.uber.org/multierr"

	"github.com/gotd/td/bin"
	tdcrypto "github.com/gotd/td/crypto"
	"github.com/gotd/td/proto/codec"
	"github.com/gotd/td/transport"
)

const maxTransportMessageSize = 1 << 24
const quickAckResponseFlag = uint32(1 << 31)

type transportListener interface {
	Accept() (transport.Conn, error)
	Close() error
	Addr() net.Addr
}

type quickAckTransport interface {
	ConsumeQuickAckRequested() bool
	SendQuickAck(ctx context.Context, token uint32) error
}

type compatTransportListener struct {
	codec    func() transport.Codec
	listener net.Listener
}

func newCompatTransportListener(codec func() transport.Codec, listener net.Listener) transportListener {
	if codec != nil {
		return transport.ListenCodec(codec, listener)
	}
	return &compatTransportListener{listener: listener}
}

// singleConnListener 是一个只产出一条「已接受」连接、随后阻塞到关闭的 net.Listener。
// 它让单条裸连接可以走 listener 形态的去混淆/codec 管线（ObfuscatedListener +
// compatTransportListener），从而把这部分阻塞读取从 accept 循环挪到每连接 goroutine。
type singleConnListener struct {
	addr net.Addr
	ch   chan net.Conn
	once sync.Once
}

func newSingleConnListener(c net.Conn) *singleConnListener {
	ch := make(chan net.Conn, 1)
	ch <- c
	return &singleConnListener{addr: c.LocalAddr(), ch: ch}
}

func (l *singleConnListener) Accept() (net.Conn, error) {
	c, ok := <-l.ch
	if !ok {
		return nil, net.ErrClosed
	}
	return c, nil
}

func (l *singleConnListener) Close() error {
	l.once.Do(func() { close(l.ch) })
	return nil
}

func (l *singleConnListener) Addr() net.Addr {
	return l.addr
}

func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
	conn, err := l.listener.Accept()
	if err != nil {
		return nil, err
	}
	defer func() {
		if rErr != nil {
			multierr.AppendInto(&rErr, conn.Close())
		}
	}()

	connCodec, reader, err := detectCompatCodec(conn)
	if err != nil {
		return nil, errors.Wrap(err, "detect codec")
	}

	return &compatTransportConn{
		conn: wrappedCompatConn{
			reader: reader,
			Conn:   conn,
		},
		codec: connCodec,
	}, nil
}

func (l *compatTransportListener) Close() error {
	return l.listener.Close()
}

func (l *compatTransportListener) Addr() net.Addr {
	return l.listener.Addr()
}

type wrappedCompatConn struct {
	reader io.Reader
	net.Conn
}

func (w wrappedCompatConn) Read(p []byte) (int, error) {
	return w.reader.Read(p)
}

type compatTransportConn struct {
	conn  net.Conn
	codec transport.Codec

	readMux  sync.Mutex
	writeMux sync.Mutex
}

func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error {
	c.writeMux.Lock()
	defer c.writeMux.Unlock()

	if err := c.conn.SetWriteDeadline(time.Time{}); err != nil {
		return errors.Wrap(err, "reset write deadline")
	}
	if deadline, ok := ctx.Deadline(); ok {
		if err := c.conn.SetWriteDeadline(deadline); err != nil {
			return errors.Wrap(err, "set write deadline")
		}
	}

	if err := c.codec.Write(c.conn, b); err != nil {
		return errors.Wrap(err, "write")
	}
	return nil
}

func (c *compatTransportConn) ConsumeQuickAckRequested() bool {
	q, ok := c.codec.(quickAckCodec)
	if !ok {
		return false
	}
	return q.consumeQuickAckRequested()
}

func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) error {
	q, ok := c.codec.(quickAckCodec)
	if !ok {
		return nil
	}

	c.writeMux.Lock()
	defer c.writeMux.Unlock()

	if err := c.conn.SetWriteDeadline(time.Time{}); err != nil {
		return errors.Wrap(err, "reset write deadline")
	}
	if deadline, ok := ctx.Deadline(); ok {
		if err := c.conn.SetWriteDeadline(deadline); err != nil {
			return errors.Wrap(err, "set write deadline")
		}
	}

	raw := q.quickAckResponse(token)
	if err := writeAll(c.conn, raw[:]); err != nil {
		return errors.Wrap(err, "write quick ack")
	}
	return nil
}

func (c *compatTransportConn) Recv(ctx context.Context, b *bin.Buffer) error {
	c.readMux.Lock()
	defer c.readMux.Unlock()

	if err := c.conn.SetReadDeadline(time.Time{}); err != nil {
		return errors.Wrap(err, "reset read deadline")
	}
	if deadline, ok := ctx.Deadline(); ok {
		if err := c.conn.SetReadDeadline(deadline); err != nil {
			return errors.Wrap(err, "set read deadline")
		}
	}

	if err := c.codec.Read(c.conn, b); err != nil {
		return errors.Wrap(err, "read")
	}
	return nil
}

func (c *compatTransportConn) Close() error {
	return c.conn.Close()
}

func detectCompatCodec(c io.Reader) (transport.Codec, io.Reader, error) {
	var buf [4]byte
	if _, err := io.ReadFull(c, buf[:1]); err != nil {
		return nil, nil, errors.Wrap(err, "read first byte")
	}

	if buf[0] == codec.AbridgedClientStart[0] {
		return &quickAckAbridgedCodec{}, c, nil
	}

	if _, err := io.ReadFull(c, buf[1:4]); err != nil {
		return nil, nil, errors.Wrap(err, "read header")
	}
	switch buf {
	case codec.IntermediateClientStart:
		return &quickAckIntermediateCodec{}, c, nil
	case codec.PaddedIntermediateClientStart:
		return &quickAckPaddedIntermediateCodec{}, c, nil
	default:
		return transport.Full.Codec(), io.MultiReader(bytes.NewReader(buf[:]), c), nil
	}
}

type quickAckCodec interface {
	transport.Codec
	consumeQuickAckRequested() bool
	quickAckResponse(token uint32) [4]byte
}

type quickAckAbridgedCodec struct {
	quickAckRequested bool
}

func (*quickAckAbridgedCodec) WriteHeader(w io.Writer) error {
	return (codec.Abridged{}).WriteHeader(w)
}

func (*quickAckAbridgedCodec) ReadHeader(r io.Reader) error {
	return (codec.Abridged{}).ReadHeader(r)
}

func (*quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error {
	if err := validateOutgoingCompatMessage(b); err != nil {
		return err
	}

	words := b.Len() >> 2
	var header [4]byte
	headerLen := 1
	if words < 0x7f {
		header[0] = byte(words)
	} else {
		header[0] = 0x7f
		header[1] = byte(words)
		header[2] = byte(words >> 8)
		header[3] = byte(words >> 16)
		headerLen = 4
	}
	return writeCompatPacket(w, header[:headerLen], b.Raw())
}

func (q *quickAckAbridgedCodec) Read(r io.Reader, b *bin.Buffer) error {
	requested, err := readQuickAckAbridged(r, b)
	if err != nil {
		return errors.Wrap(err, "read abridged")
	}
	q.quickAckRequested = requested
	return checkCompatProtocolError(b)
}

func (q *quickAckAbridgedCodec) consumeQuickAckRequested() bool {
	v := q.quickAckRequested
	q.quickAckRequested = false
	return v
}

func (*quickAckAbridgedCodec) quickAckResponse(token uint32) [4]byte {
	var raw [4]byte
	binary.BigEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
	return raw
}

type quickAckIntermediateCodec struct {
	quickAckRequested bool
}

func (*quickAckIntermediateCodec) WriteHeader(w io.Writer) error {
	return (codec.Intermediate{}).WriteHeader(w)
}

func (*quickAckIntermediateCodec) ReadHeader(r io.Reader) error {
	return (codec.Intermediate{}).ReadHeader(r)
}

func (*quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
	if err := validateOutgoingCompatMessage(b); err != nil {
		return err
	}
	var header [4]byte
	binary.LittleEndian.PutUint32(header[:], uint32(b.Len()))
	return writeCompatPacket(w, header[:], b.Raw())
}

func (q *quickAckIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
	requested, err := readQuickAckIntermediate(r, b, false)
	if err != nil {
		return errors.Wrap(err, "read intermediate")
	}
	q.quickAckRequested = requested
	return checkCompatProtocolError(b)
}

func (q *quickAckIntermediateCodec) consumeQuickAckRequested() bool {
	v := q.quickAckRequested
	q.quickAckRequested = false
	return v
}

func (*quickAckIntermediateCodec) quickAckResponse(token uint32) [4]byte {
	var raw [4]byte
	binary.LittleEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
	return raw
}

type quickAckPaddedIntermediateCodec struct {
	quickAckRequested bool
}

func (*quickAckPaddedIntermediateCodec) WriteHeader(w io.Writer) error {
	return (codec.PaddedIntermediate{}).WriteHeader(w)
}

func (*quickAckPaddedIntermediateCodec) ReadHeader(r io.Reader) error {
	return (codec.PaddedIntermediate{}).ReadHeader(r)
}

func (*quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
	if err := validateOutgoingCompatMessage(b); err != nil {
		return err
	}
	var padding [4]byte
	if _, err := io.ReadFull(tdcrypto.DefaultRand(), padding[:]); err != nil {
		return err
	}
	n := int(padding[0] % 4)
	payload := b.Raw()
	if n > 0 {
		payload = append(append([]byte(nil), payload...), padding[:n]...)
	}
	var header [4]byte
	binary.LittleEndian.PutUint32(header[:], uint32(len(payload)))
	return writeCompatPacket(w, header[:], payload)
}

func (q *quickAckPaddedIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
	requested, err := readQuickAckIntermediate(r, b, true)
	if err != nil {
		return errors.Wrap(err, "read padded intermediate")
	}
	q.quickAckRequested = requested
	return checkCompatProtocolError(b)
}

func (q *quickAckPaddedIntermediateCodec) consumeQuickAckRequested() bool {
	v := q.quickAckRequested
	q.quickAckRequested = false
	return v
}

func (*quickAckPaddedIntermediateCodec) quickAckResponse(token uint32) [4]byte {
	var raw [4]byte
	binary.LittleEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
	return raw
}

func readQuickAckAbridged(r io.Reader, b *bin.Buffer) (bool, error) {
	var first [1]byte
	if _, err := io.ReadFull(r, first[:]); err != nil {
		return false, err
	}

	requested := first[0]&0x80 != 0
	lengthByte := first[0] & 0x7f
	var n int
	if lengthByte == 0x7f {
		var tail [3]byte
		if _, err := io.ReadFull(r, tail[:]); err != nil {
			return false, err
		}
		words := uint32(tail[0]) | uint32(tail[1])<<8 | uint32(tail[2])<<16
		n = int(words << 2)
	} else {
		n = int(lengthByte) << 2
	}

	if err := validateCompatTransportLength(n); err != nil {
		return false, err
	}
	resetCompatBufferN(b, n)
	if _, err := io.ReadFull(r, b.Buf); err != nil {
		return false, errors.Wrap(err, "read payload")
	}
	return requested, nil
}

func readQuickAckIntermediate(r io.Reader, b *bin.Buffer, padding bool) (bool, error) {
	var lengthBuf [4]byte
	if _, err := io.ReadFull(r, lengthBuf[:]); err != nil {
		return false, errors.Wrap(err, "read length")
	}
	rawLength := binary.LittleEndian.Uint32(lengthBuf[:])
	requested := rawLength&quickAckResponseFlag != 0
	n := int(rawLength &^ quickAckResponseFlag)
	if err := validateCompatTransportLength(n); err != nil {
		return false, err
	}
	resetCompatBufferN(b, n)
	if _, err := io.ReadFull(r, b.Buf); err != nil {
		return false, errors.Wrap(err, "read payload")
	}
	if padding {
		paddingLength := n % 4
		b.Buf = b.Buf[:n-paddingLength]
	}
	return requested, nil
}

func validateOutgoingCompatMessage(b *bin.Buffer) error {
	n := b.Len()
	if err := validateCompatTransportLength(n); err != nil {
		return err
	}
	if n%bin.Word != 0 {
		return fmt.Errorf("invalid message length %d: not aligned to %d", n, bin.Word)
	}
	return nil
}

func writeCompatPacket(w io.Writer, header, payload []byte) error {
	packet := make([]byte, 0, len(header)+len(payload))
	packet = append(packet, header...)
	packet = append(packet, payload...)
	return writeAll(w, packet)
}

func writeAll(w io.Writer, p []byte) error {
	for len(p) > 0 {
		n, err := w.Write(p)
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrShortWrite
		}
		p = p[n:]
	}
	return nil
}

func validateCompatTransportLength(n int) error {
	if n <= 0 || n > maxTransportMessageSize {
		return fmt.Errorf("invalid message length %d", n)
	}
	return nil
}

func resetCompatBufferN(b *bin.Buffer, n int) {
	if cap(b.Buf) < n {
		b.Buf = make([]byte, n)
		return
	}
	b.Buf = b.Buf[:n]
}

func checkCompatProtocolError(b *bin.Buffer) error {
	if b.Len() != bin.Word {
		return nil
	}
	code, err := b.Int32()
	if err != nil {
		return err
	}
	return &codec.ProtocolErr{Code: -code}
}