package store

import (
	"maps"
	"net/netip"
	"slices"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/netbirdio/netbird/client/internal/netflow/types"
)

func NewMemoryStore() *Memory {
	return &Memory{
		events: make(map[uuid.UUID]*types.Event),
	}
}

type Memory struct {
	mux    sync.Mutex
	events map[uuid.UUID]*types.Event
}

type AggregatingMemory struct {
	Memory
}

func (m *Memory) StoreEvent(event *types.Event) {
	m.mux.Lock()
	defer m.mux.Unlock()
	m.events[event.ID] = event
}

func (m *Memory) Close() {
	m.mux.Lock()
	defer m.mux.Unlock()
	clear(m.events)
}

func (m *Memory) GetEvents() []*types.Event {
	m.mux.Lock()
	defer m.mux.Unlock()
	events := make([]*types.Event, 0, len(m.events))
	for _, event := range m.events {
		events = append(events, event)
	}
	return events
}

func (m *Memory) DeleteEvents(ids []uuid.UUID) {
	m.mux.Lock()
	defer m.mux.Unlock()
	for _, id := range ids {
		delete(m.events, id)
	}
}

func NewAggregatingMemoryStore() *AggregatingMemory {
	return &AggregatingMemory{Memory{events: make(map[uuid.UUID]*types.Event)}}
}

func (am *AggregatingMemory) ResetAggregationWindow() types.FlowEventAggregator {
	am.mux.Lock()
	defer am.mux.Unlock()

	toret := AggregatingMemory{Memory: Memory{events: am.events}}
	am.events = make(map[uuid.UUID]*types.Event)

	return &toret
}

type aggregationKey struct {
	destAddr netip.Addr
	destPort uint16
	protocol uint8
	icmpType uint8
	unique   int64 // used to prevent aggregation on non icmp/udp/tcp events
}

func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event {
	aggregated := make(map[aggregationKey]*types.Event)
	for _, v := range am.events {
		lookupKey := aggregationKey{destAddr: v.DestIP, destPort: v.DestPort, protocol: uint8(v.Protocol), icmpType: v.ICMPType}
		if _, ok := aggregated[lookupKey]; !ok {
			aggregated[lookupKey] = v.Clone()
			event := aggregated[lookupKey]

			if event.Protocol != types.ICMP && event.Protocol != types.ICMPv6 && event.Protocol != types.UDP && event.Protocol != types.TCP {
				lookupKey.unique = time.Now().UnixNano() // to make the lookup key unique so we don't aggregate on it
				continue
			}

			switch event.Type {
			case types.TypeStart:
				event.NumOfStarts += 1
			case types.TypeDrop:
				event.NumOfDrops += 1
			case types.TypeEnd:
				event.NumOfEnds += 1
			}
			continue
		}

		aggregatedEvent := aggregated[lookupKey]
		if aggregatedEvent.Protocol != types.ICMP && aggregatedEvent.Protocol != types.ICMPv6 && aggregatedEvent.Protocol != types.UDP && aggregatedEvent.Protocol != types.TCP {
			continue // we don't aggregate this type of events; shouldn't ever get here
		}

		// track the number of connections, duration?, open and close events?
		aggregatedEvent.RxBytes += v.RxBytes
		aggregatedEvent.RxPackets += v.RxPackets
		aggregatedEvent.TxBytes += v.TxBytes
		aggregatedEvent.TxPackets += v.TxPackets
		switch v.Type {
		case types.TypeStart:
			aggregatedEvent.NumOfStarts += 1
		case types.TypeDrop:
			aggregatedEvent.NumOfDrops += 1
		case types.TypeEnd:
			aggregatedEvent.NumOfEnds += 1
		}
		if aggregatedEvent.Timestamp.Compare(v.Timestamp) > 0 {
			aggregatedEvent.Timestamp = v.Timestamp
			aggregatedEvent.ID = v.ID
			aggregatedEvent.Type = v.Type
		}
		// do we aggregate icmp by code?
	}

	return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here
}