package engine

import (
	"errors"
	"net"
	"sync"

	"github.com/apernet/OpenGFW/analyzer"
	"github.com/apernet/OpenGFW/io"
	"github.com/apernet/OpenGFW/modifier"
	"github.com/apernet/OpenGFW/ruleset"

	"github.com/bwmarrin/snowflake"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	lru "github.com/hashicorp/golang-lru/v2"
)

// udpVerdict is a subset of io.Verdict for UDP streams.
// For UDP, we support all verdicts.
type udpVerdict io.Verdict

const (
	udpVerdictAccept       = udpVerdict(io.VerdictAccept)
	udpVerdictAcceptModify = udpVerdict(io.VerdictAcceptModify)
	udpVerdictAcceptStream = udpVerdict(io.VerdictAcceptStream)
	udpVerdictDrop         = udpVerdict(io.VerdictDrop)
	udpVerdictDropStream   = udpVerdict(io.VerdictDropStream)
)

var errInvalidModifier = errors.New("invalid modifier")

type udpContext struct {
	Verdict udpVerdict
	Packet  []byte
}

type udpStreamFactory struct {
	WorkerID int
	Logger   Logger
	Node     *snowflake.Node

	RulesetMutex sync.RWMutex
	Ruleset      ruleset.Ruleset
}

func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, uc *udpContext) *udpStream {
	id := f.Node.Generate()
	ipSrc, ipDst := net.IP(ipFlow.Src().Raw()), net.IP(ipFlow.Dst().Raw())
	info := ruleset.StreamInfo{
		ID:       id.Int64(),
		Protocol: ruleset.ProtocolUDP,
		SrcIP:    ipSrc,
		DstIP:    ipDst,
		SrcPort:  uint16(udp.SrcPort),
		DstPort:  uint16(udp.DstPort),
		Props:    make(analyzer.CombinedPropMap),
	}
	f.Logger.UDPStreamNew(f.WorkerID, info)
	f.RulesetMutex.RLock()
	rs := f.Ruleset
	f.RulesetMutex.RUnlock()
	ans := analyzersToUDPAnalyzers(rs.Analyzers(info))
	if len(ans) == 0 {
		uc.Verdict = udpVerdictAcceptStream
		f.Logger.UDPStreamAction(info, ruleset.ActionAllow, true)
		// a udpStream with no activeEntries is a no-op
		return &udpStream{}
	}
	// Create entries for each analyzer
	entries := make([]*udpStreamEntry, 0, len(ans))
	for _, a := range ans {
		entries = append(entries, &udpStreamEntry{
			Name: a.Name(),
			Stream: a.NewUDP(analyzer.UDPInfo{
				SrcIP:   ipSrc,
				DstIP:   ipDst,
				SrcPort: uint16(udp.SrcPort),
				DstPort: uint16(udp.DstPort),
			}, &analyzerLogger{
				StreamID: id.Int64(),
				Name:     a.Name(),
				Logger:   f.Logger,
			}),
			HasLimit: a.Limit() > 0,
			Quota:    a.Limit(),
		})
	}
	return &udpStream{
		info:          info,
		virgin:        true,
		logger:        f.Logger,
		ruleset:       rs,
		activeEntries: entries,
	}
}

func (f *udpStreamFactory) UpdateRuleset(r ruleset.Ruleset) error {
	f.RulesetMutex.Lock()
	defer f.RulesetMutex.Unlock()
	f.Ruleset = r
	return nil
}

type udpStreamManager struct {
	factory *udpStreamFactory
	streams *lru.Cache[uint32, *udpStreamValue]
}

type udpStreamValue struct {
	Stream  *udpStream
	IPFlow  gopacket.Flow
	UDPFlow gopacket.Flow
}

func (v *udpStreamValue) Match(ipFlow, udpFlow gopacket.Flow) (ok, rev bool) {
	fwd := v.IPFlow == ipFlow && v.UDPFlow == udpFlow
	rev = v.IPFlow == ipFlow.Reverse() && v.UDPFlow == udpFlow.Reverse()
	return fwd || rev, rev
}

func newUDPStreamManager(factory *udpStreamFactory, maxStreams int) (*udpStreamManager, error) {
	ss, err := lru.New[uint32, *udpStreamValue](maxStreams)
	if err != nil {
		return nil, err
	}
	return &udpStreamManager{
		factory: factory,
		streams: ss,
	}, nil
}

func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flow, udp *layers.UDP, uc *udpContext) {
	rev := false
	value, ok := m.streams.Get(streamID)
	if !ok {
		// New stream
		value = &udpStreamValue{
			Stream:  m.factory.New(ipFlow, udp.TransportFlow(), udp, uc),
			IPFlow:  ipFlow,
			UDPFlow: udp.TransportFlow(),
		}
		m.streams.Add(streamID, value)
	} else {
		// Stream ID exists, but is it really the same stream?
		ok, rev = value.Match(ipFlow, udp.TransportFlow())
		if !ok {
			// It's not - close the old stream & replace it with a new one
			value.Stream.Close()
			value = &udpStreamValue{
				Stream:  m.factory.New(ipFlow, udp.TransportFlow(), udp, uc),
				IPFlow:  ipFlow,
				UDPFlow: udp.TransportFlow(),
			}
			m.streams.Add(streamID, value)
		}
	}
	if value.Stream.Accept(udp, rev, uc) {
		value.Stream.Feed(udp, rev, uc)
	}
}

type udpStream struct {
	info          ruleset.StreamInfo
	virgin        bool // true if no packets have been processed
	logger        Logger
	ruleset       ruleset.Ruleset
	activeEntries []*udpStreamEntry
	doneEntries   []*udpStreamEntry
}

type udpStreamEntry struct {
	Name     string
	Stream   analyzer.UDPStream
	HasLimit bool
	Quota    int
}

func (s *udpStream) Accept(udp *layers.UDP, rev bool, uc *udpContext) bool {
	// Only accept packets if we still have active entries
	return len(s.activeEntries) > 0
}

func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
	updated := false
	for i := len(s.activeEntries) - 1; i >= 0; i-- {
		// Important: reverse order so we can remove entries
		entry := s.activeEntries[i]
		update, closeUpdate, done := s.feedEntry(entry, rev, udp.Payload)
		updated = updated || processPropUpdate(s.info.Props, entry.Name, update)
		updated = updated || processPropUpdate(s.info.Props, entry.Name, closeUpdate)
		if done {
			s.activeEntries = append(s.activeEntries[:i], s.activeEntries[i+1:]...)
			s.doneEntries = append(s.doneEntries, entry)
		}
	}
	if updated || s.virgin {
		s.virgin = false
		s.logger.UDPStreamPropUpdate(s.info, false)
		// Match properties against ruleset
		result, err := s.ruleset.Match(s.info)
		if err != nil {
			s.logger.MatchError(s.info, err)
		}
		action := result.Action
		if action == ruleset.ActionModify {
			// Call the modifier instance
			udpMI, ok := result.ModInstance.(modifier.UDPModifierInstance)
			if !ok {
				// Not for UDP, fallback to maybe
				s.logger.ModifyError(s.info, errInvalidModifier)
				action = ruleset.ActionMaybe
			} else {
				uc.Packet, err = udpMI.Process(udp.Payload)
				if err != nil {
					// Modifier error, fallback to maybe
					s.logger.ModifyError(s.info, err)
					action = ruleset.ActionMaybe
				}
			}
		}
		if action != ruleset.ActionMaybe {
			var final bool
			uc.Verdict, final = actionToUDPVerdict(action)
			s.logger.UDPStreamAction(s.info, action, false)
			if final {
				s.closeActiveEntries()
			}
		}
	}
	if len(s.activeEntries) == 0 && uc.Verdict == udpVerdictAccept {
		// All entries are done but no verdict issued, accept stream
		uc.Verdict = udpVerdictAcceptStream
		s.logger.UDPStreamAction(s.info, ruleset.ActionAllow, true)
	}
}

func (s *udpStream) Close() {
	s.closeActiveEntries()
}

func (s *udpStream) closeActiveEntries() {
	// Signal close to all active entries & move them to doneEntries
	updated := false
	for _, entry := range s.activeEntries {
		update := entry.Stream.Close(false)
		updated = updated || processPropUpdate(s.info.Props, entry.Name, update)
	}
	if updated {
		s.logger.UDPStreamPropUpdate(s.info, true)
	}
	s.doneEntries = append(s.doneEntries, s.activeEntries...)
	s.activeEntries = nil
}

func (s *udpStream) feedEntry(entry *udpStreamEntry, rev bool, data []byte) (update *analyzer.PropUpdate, closeUpdate *analyzer.PropUpdate, done bool) {
	update, done = entry.Stream.Feed(rev, data)
	if entry.HasLimit {
		entry.Quota -= len(data)
		if entry.Quota <= 0 {
			// Quota exhausted, signal close & move to doneEntries
			closeUpdate = entry.Stream.Close(true)
			done = true
		}
	}
	return
}

func analyzersToUDPAnalyzers(ans []analyzer.Analyzer) []analyzer.UDPAnalyzer {
	udpAns := make([]analyzer.UDPAnalyzer, 0, len(ans))
	for _, a := range ans {
		if udpM, ok := a.(analyzer.UDPAnalyzer); ok {
			udpAns = append(udpAns, udpM)
		}
	}
	return udpAns
}

func actionToUDPVerdict(a ruleset.Action) (v udpVerdict, final bool) {
	switch a {
	case ruleset.ActionMaybe:
		return udpVerdictAccept, false
	case ruleset.ActionAllow:
		return udpVerdictAcceptStream, true
	case ruleset.ActionBlock:
		return udpVerdictDropStream, true
	case ruleset.ActionDrop:
		return udpVerdictDrop, false
	case ruleset.ActionModify:
		return udpVerdictAcceptModify, false
	default:
		// Should never happen
		return udpVerdictAccept, false
	}
}