package udp

import (
	"bytes"
	"errors"
	"sort"

	"git.difuse.io/Difuse/Mellaris/analyzer"
	"git.difuse.io/Difuse/Mellaris/analyzer/internal"
	"git.difuse.io/Difuse/Mellaris/analyzer/udp/internal/quic"
	"git.difuse.io/Difuse/Mellaris/analyzer/utils"
)

const (
	quicInvalidCountThreshold = 16
	quicMaxCryptoDataLen      = 256 * 1024
)

var (
	_ analyzer.UDPAnalyzer = (*QUICAnalyzer)(nil)
	_ analyzer.UDPStream   = (*quicStream)(nil)
)

type QUICAnalyzer struct{}

func (a *QUICAnalyzer) Name() string {
	return "quic"
}

func (a *QUICAnalyzer) Limit() int {
	return 0
}

func (a *QUICAnalyzer) NewUDP(info analyzer.UDPInfo, logger analyzer.Logger) analyzer.UDPStream {
	return &quicStream{
		logger: logger,
		frames: make(map[int64][]byte),
	}
}

type quicStream struct {
	logger       analyzer.Logger
	invalidCount int
	debugCount   int
	frames       map[int64][]byte
	maxEnd       int64
	connIDs      [][]byte
	lastHint     quic.DecryptSuccessHint
	hasLastHint  bool
}

func (s *quicStream) Feed(rev bool, data []byte) (u *analyzer.PropUpdate, done bool) {
	// minimal data size: protocol version (2 bytes) + random (32 bytes) +
	//   + session ID (1 byte) + cipher suites (4 bytes) +
	//   + compression methods (2 bytes) + no extensions
	const minDataSize = 41

	var parsedHdr *quic.Header
	if hdr, _, err := quic.ParseInitialHeader(data); err == nil {
		parsedHdr = hdr
		s.rememberConnID(hdr.DestConnectionID)
		s.rememberConnID(hdr.SrcConnectionID)
	}
	hint := quic.DecryptSuccessHint{}
	opts := &quic.ReadCryptoFramesOptions{
		AdditionalConnectionIDs: s.connIDs,
		TryServerSecret:         true,
		SuccessHint:             &hint,
	}
	if s.hasLastHint {
		opts.PreferredConnectionID = append([]byte(nil), s.lastHint.ConnectionID...)
		opts.PreferredSecretLabel = s.lastHint.SecretLabel
		opts.PreferredPNMax = s.lastHint.PacketNumberMax
		opts.HasPreferredPNMax = true
	}
	frs, err := quic.ReadCryptoFramesWithOptions(data, opts)
	if err != nil {
		if errors.Is(err, quic.ErrNotInitialPacket) {
			return nil, false
		}
		if s.debugCount < 4 {
			if parsedHdr != nil {
				s.logger.Debugf(
					"failed to read QUIC CRYPTO frames: version=%x dcid_len=%d scid_len=%d token_len=%d conn_id_hints=%d err=%v",
					parsedHdr.Version,
					len(parsedHdr.DestConnectionID),
					len(parsedHdr.SrcConnectionID),
					len(parsedHdr.Token),
					len(s.connIDs),
					err,
				)
			} else {
				s.logger.Debugf("failed to read QUIC CRYPTO frames: %v", err)
			}
			s.debugCount++
		}
		s.invalidCount++
		return nil, s.invalidCount >= quicInvalidCountThreshold
	}
	if len(hint.ConnectionID) > 0 {
		s.lastHint = hint
		s.hasLastHint = true
		s.rememberConnID(hint.ConnectionID)
	}
	if len(frs) == 0 {
		s.invalidCount++
		return nil, s.invalidCount >= quicInvalidCountThreshold
	}
	s.invalidCount = 0
	for _, f := range frs {
		s.mergeFrame(f.Offset, f.Data)
	}
	pl := s.contiguousPayloadFromZero()
	if len(pl) < 4 {
		return nil, false
	}

	if pl[0] != internal.TypeClientHello {
		// Not a ClientHello (e.g. server-direction CRYPTO); ignore.
		if s.debugCount < 4 {
			s.logger.Debugf("CRYPTO payload does not start with ClientHello: type=%d", pl[0])
			s.debugCount++
		}
		return nil, false
	}

	chLen := int(pl[1])<<16 | int(pl[2])<<8 | int(pl[3])
	if chLen < minDataSize {
		s.invalidCount++
		return nil, s.invalidCount >= quicInvalidCountThreshold
	}
	if len(pl) < 4+chLen {
		// Wait for more CRYPTO data from subsequent packets.
		return nil, false
	}

	m := internal.ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: pl[4 : 4+chLen]})
	if m == nil {
		if s.debugCount < 4 {
			s.logger.Debugf("failed to parse TLS ClientHello from QUIC CRYPTO payload")
			s.debugCount++
		}
		s.invalidCount++
		return nil, s.invalidCount >= quicInvalidCountThreshold
	}

	return &analyzer.PropUpdate{
		Type: analyzer.PropUpdateMerge,
		M:    analyzer.PropMap{"req": m},
	}, true
}

func (s *quicStream) Close(limited bool) *analyzer.PropUpdate {
	return nil
}

func (s *quicStream) mergeFrame(offset int64, data []byte) {
	if len(data) == 0 || offset < 0 {
		return
	}
	if s.frames == nil {
		s.frames = make(map[int64][]byte)
	}
	if _, exists := s.frames[offset]; exists {
		return
	}
	s.frames[offset] = append([]byte(nil), data...)
	end := offset + int64(len(data))
	if end > s.maxEnd {
		s.maxEnd = end
	}
}

func (s *quicStream) contiguousPayloadFromZero() []byte {
	if len(s.frames) == 0 || s.maxEnd <= 0 || s.maxEnd > quicMaxCryptoDataLen {
		return nil
	}
	keys := make([]int64, 0, len(s.frames))
	for k := range s.frames {
		keys = append(keys, k)
	}
	sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
	if keys[0] != 0 {
		return nil
	}
	out := make([]byte, 0, s.maxEnd)
	next := int64(0)
	for _, k := range keys {
		if k > next {
			break
		}
		frame := s.frames[k]
		frameEnd := k + int64(len(frame))
		if frameEnd <= next {
			continue
		}
		start := next - k
		out = append(out, frame[start:]...)
		next = frameEnd
		if next >= quicMaxCryptoDataLen {
			break
		}
	}
	return out
}

func (s *quicStream) rememberConnID(cid []byte) {
	if len(cid) == 0 {
		return
	}
	for _, existing := range s.connIDs {
		if bytes.Equal(existing, cid) {
			return
		}
	}
	// Keep this bounded; QUIC CID rotations are small for handshake parsing.
	if len(s.connIDs) >= 8 {
		s.connIDs = s.connIDs[1:]
	}
	s.connIDs = append(s.connIDs, append([]byte(nil), cid...))
}