package udp

import (
	"encoding/binary"
	"eq2emu/internal/opcodes"
	"errors"
	"fmt"
	"sort"
)

// FragmentManager handles packet fragmentation and reassembly
type FragmentManager struct {
	fragments map[uint16]*FragmentGroup // Active fragment groups by base sequence
	maxLength uint32                    // Maximum packet size before fragmentation
}

// FragmentGroup tracks fragments belonging to the same original packet
type FragmentGroup struct {
	BaseSequence uint16          // Base sequence number
	TotalLength  uint32          // Expected total reassembled length
	Fragments    []FragmentPiece // Individual fragment pieces
	FirstSeen    bool            // Whether we've seen the first fragment
}

// FragmentPiece represents a single fragment
type FragmentPiece struct {
	Sequence uint16 // Fragment sequence number
	Data     []byte // Fragment payload data
	IsFirst  bool   // Whether this is the first fragment
}

// NewFragmentManager creates a manager with specified maximum packet length
func NewFragmentManager(maxLength uint32) *FragmentManager {
	return &FragmentManager{
		fragments: make(map[uint16]*FragmentGroup),
		maxLength: maxLength,
	}
}

// FragmentPacket splits large packets into fragments
// Returns nil if packet doesn't need fragmentation
func (fm *FragmentManager) FragmentPacket(data []byte, startSeq uint16) []*ProtocolPacket {
	if uint32(len(data)) <= fm.maxLength {
		return nil // No fragmentation needed
	}

	totalLength := uint32(len(data))
	chunkSize := int(fm.maxLength - 6) // Reserve 6 bytes for headers
	if chunkSize <= 0 {
		chunkSize = 1
	}

	var packets []*ProtocolPacket
	seq := startSeq

	for offset := 0; offset < len(data); offset += chunkSize {
		end := offset + chunkSize
		if end > len(data) {
			end = len(data)
		}

		var fragmentData []byte
		if offset == 0 {
			// First fragment includes total length
			fragmentData = make([]byte, 6+end-offset)
			binary.BigEndian.PutUint16(fragmentData[0:2], seq)
			binary.LittleEndian.PutUint32(fragmentData[2:6], totalLength)
			copy(fragmentData[6:], data[offset:end])
		} else {
			// Subsequent fragments
			fragmentData = make([]byte, 2+end-offset)
			binary.BigEndian.PutUint16(fragmentData[0:2], seq)
			copy(fragmentData[2:], data[offset:end])
		}

		packets = append(packets, &ProtocolPacket{
			Opcode: opcodes.OpFragment,
			Data:   fragmentData,
		})
		seq++
	}

	return packets
}

// ProcessFragment handles incoming fragments and returns complete packet when ready
func (fm *FragmentManager) ProcessFragment(packet *ProtocolPacket) ([]byte, bool, error) {
	if len(packet.Data) < 2 {
		return nil, false, errors.New("fragment too small")
	}

	seq := binary.BigEndian.Uint16(packet.Data[0:2])

	// Parse fragment data
	fragment := FragmentPiece{Sequence: seq}

	if len(packet.Data) >= 6 {
		// Check if this is the first fragment (has total length)
		possibleLength := binary.LittleEndian.Uint32(packet.Data[2:6])
		if possibleLength > 0 && possibleLength < 10*1024*1024 { // Reasonable limit
			fragment.IsFirst = true
			fragment.Data = packet.Data[6:]

			// Create new fragment group
			group := &FragmentGroup{
				BaseSequence: seq,
				TotalLength:  possibleLength,
				Fragments:    []FragmentPiece{fragment},
				FirstSeen:    true,
			}
			fm.fragments[seq] = group

			return fm.tryAssemble(seq)
		}
	}

	// Not first fragment - find matching group
	fragment.Data = packet.Data[2:]
	group := fm.findFragmentGroup(seq)
	if group == nil {
		return nil, false, errors.New("orphaned fragment")
	}

	group.Fragments = append(group.Fragments, fragment)
	return fm.tryAssemble(group.BaseSequence)
}

// findFragmentGroup locates the fragment group for a sequence number
func (fm *FragmentManager) findFragmentGroup(seq uint16) *FragmentGroup {
	// Look for groups where this sequence fits
	for baseSeq, group := range fm.fragments {
		if seq >= baseSeq && seq < baseSeq+100 { // Reasonable window
			return group
		}
	}
	return nil
}

// tryAssemble attempts to reassemble fragments into complete packet
func (fm *FragmentManager) tryAssemble(baseSeq uint16) ([]byte, bool, error) {
	group, exists := fm.fragments[baseSeq]
	if !exists || !group.FirstSeen {
		return nil, false, nil
	}

	// Calculate expected fragment count
	chunkSize := int(fm.maxLength - 6)
	expectedCount := int(group.TotalLength) / chunkSize
	if int(group.TotalLength)%chunkSize != 0 {
		expectedCount++
	}

	if len(group.Fragments) < expectedCount {
		return nil, false, nil // Still waiting for fragments
	}

	// Sort fragments by sequence number
	sort.Slice(group.Fragments, func(i, j int) bool {
		return group.Fragments[i].Sequence < group.Fragments[j].Sequence
	})

	// Reassemble packet
	result := make([]byte, 0, group.TotalLength)
	for _, frag := range group.Fragments[:expectedCount] {
		result = append(result, frag.Data...)
	}

	// Validate length
	if uint32(len(result)) != group.TotalLength {
		delete(fm.fragments, baseSeq)
		return nil, false, fmt.Errorf("assembled length %d != expected %d", len(result), group.TotalLength)
	}

	// Clean up
	delete(fm.fragments, baseSeq)
	return result, true, nil
}

// CleanupStale removes old incomplete fragment groups
func (fm *FragmentManager) CleanupStale(maxAge uint16) {
	// Simple cleanup - remove groups with very old base sequences
	for baseSeq := range fm.fragments {
		if baseSeq < maxAge {
			delete(fm.fragments, baseSeq)
		}
	}
}

// GetStats returns fragmentation statistics
func (fm *FragmentManager) GetStats() FragmentStats {
	return FragmentStats{
		ActiveGroups: len(fm.fragments),
		MaxLength:    fm.maxLength,
	}
}

// FragmentStats contains fragmentation statistics
type FragmentStats struct {
	ActiveGroups int    // Number of incomplete fragment groups
	MaxLength    uint32 // Maximum packet length setting
}

// Clear removes all pending fragments
func (fm *FragmentManager) Clear() {
	fm.fragments = make(map[uint16]*FragmentGroup)
}

// SetMaxLength updates the maximum packet length
func (fm *FragmentManager) SetMaxLength(maxLength uint32) {
	fm.maxLength = maxLength
}