package main

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"net/netip"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/ti-mo/conntrack"
	"github.com/ti-mo/netfilter"
)

type conntrackMetrics struct {
	Emitted     atomic.Uint64
	Eligible    atomic.Uint64
	RateDropped atomic.Uint64
	SampledOut  atomic.Uint64
	Up          atomic.Bool
}

type conntrackRecord struct {
	DurationMilliseconds uint64         `json:"duration_milliseconds"`
	EndedAt              time.Time      `json:"ended_at"`
	Original             conntrackTuple `json:"original"`
	OriginalBytes        uint64         `json:"original_bytes"`
	OriginalPackets      uint64         `json:"original_packets"`
	Reply                conntrackTuple `json:"reply"`
	ReplyBytes           uint64         `json:"reply_bytes"`
	ReplyPackets         uint64         `json:"reply_packets"`
	SchemaVersion        int            `json:"schema_version"`
	Service              string         `json:"service"`
	StartedAt            time.Time      `json:"started_at"`
	TCPState             string         `json:"tcp_state"`
	Type                 string         `json:"type"`
}

type conntrackTuple struct {
	DestinationAddress string `json:"destination_address"`
	DestinationPort    uint16 `json:"destination_port"`
	SourceAddress      string `json:"source_address"`
	SourcePort         uint16 `json:"source_port"`
}

type tokenBucket struct {
	burst     float64
	last      time.Time
	rate      float64
	remaining float64
}

func newTokenBucket(rate, burst int, now time.Time) tokenBucket {
	return tokenBucket{burst: float64(burst), last: now, rate: float64(rate), remaining: float64(burst)}
}

func (bucket *tokenBucket) allow(now time.Time) bool {
	if now.After(bucket.last) {
		bucket.remaining = min(bucket.burst, bucket.remaining+now.Sub(bucket.last).Seconds()*bucket.rate)
		bucket.last = now
	}
	if bucket.remaining < 1 {
		return false
	}
	bucket.remaining--
	return true
}

func observeConntrack(ctx context.Context, configuration relayConfig, metrics *conntrackMetrics, output io.Writer) error {
	connection, err := conntrack.Dial(nil)
	if err != nil {
		return err
	}
	defer connection.Close()
	events := make(chan conntrack.Event, configuration.ConntrackEventBuffer)
	errors, err := connection.Listen(events, 1, []netfilter.NetlinkGroup{netfilter.GroupCTDestroy})
	if err != nil {
		return err
	}
	metrics.Up.Store(true)
	defer metrics.Up.Store(false)
	go func() {
		<-ctx.Done()
		connection.Close()
	}()
	encoder := json.NewEncoder(output)
	limiter := newTokenBucket(configuration.ConntrackEmissionRatePerSecond, configuration.ConntrackEmissionBurst, time.Now())
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case err := <-errors:
			return err
		case event := <-events:
			record, eligible := conntrackRecordForFlow(event.Flow, configuration.ConntrackServices)
			if !eligible {
				continue
			}
			metrics.Eligible.Add(1)
			if !sampleConntrackFlow(event.Flow, configuration.ConntrackSampleDenominator) {
				metrics.SampledOut.Add(1)
				continue
			}
			if !limiter.allow(time.Now()) {
				metrics.RateDropped.Add(1)
				continue
			}
			if err := encoder.Encode(record); err != nil {
				return err
			}
			metrics.Emitted.Add(1)
		}
	}
}

func conntrackRecordForFlow(flow *conntrack.Flow, services []conntrackService) (conntrackRecord, bool) {
	if flow == nil || flow.TupleOrig.Proto.Protocol != syscall.IPPROTO_TCP || flow.TupleReply.Proto.Protocol != syscall.IPPROTO_TCP {
		return conntrackRecord{}, false
	}
	service := ""
	for _, candidate := range services {
		if flow.TupleReply.IP.SourceAddress != candidate.TargetAddress {
			continue
		}
		for _, port := range candidate.Ports {
			if flow.TupleOrig.Proto.DestinationPort == port {
				service = candidate.Name
				break
			}
		}
	}
	if service == "" || !validTuple(flow.TupleOrig) || !validTuple(flow.TupleReply) || flow.Timestamp.Start.IsZero() {
		return conntrackRecord{}, false
	}
	endedAt := flow.Timestamp.Stop
	if endedAt.IsZero() {
		endedAt = time.Now()
	}
	duration := endedAt.Sub(flow.Timestamp.Start)
	if duration < 0 {
		return conntrackRecord{}, false
	}
	state := "unknown"
	if flow.ProtoInfo.TCP != nil {
		state = tcpState(flow.ProtoInfo.TCP.State)
	}
	return conntrackRecord{
		DurationMilliseconds: uint64(duration.Milliseconds()),
		EndedAt:              endedAt.UTC(),
		Original:             recordTuple(flow.TupleOrig),
		OriginalBytes:        flow.CountersOrig.Bytes,
		OriginalPackets:      flow.CountersOrig.Packets,
		Reply:                recordTuple(flow.TupleReply),
		ReplyBytes:           flow.CountersReply.Bytes,
		ReplyPackets:         flow.CountersReply.Packets,
		SchemaVersion:        1,
		Service:              service,
		StartedAt:            flow.Timestamp.Start.UTC(),
		TCPState:             state,
		Type:                 "oci_l4_relay_conntrack",
	}, true
}

func validTuple(tuple conntrack.Tuple) bool {
	return tuple.IP.SourceAddress.IsValid() && tuple.IP.DestinationAddress.IsValid() && tuple.Proto.SourcePort != 0 && tuple.Proto.DestinationPort != 0
}

func recordTuple(tuple conntrack.Tuple) conntrackTuple {
	return conntrackTuple{
		DestinationAddress: tuple.IP.DestinationAddress.String(),
		DestinationPort:    tuple.Proto.DestinationPort,
		SourceAddress:      tuple.IP.SourceAddress.String(),
		SourcePort:         tuple.Proto.SourcePort,
	}
}

func sampleConntrackFlow(flow *conntrack.Flow, denominator uint64) bool {
	if denominator == 1 {
		return true
	}
	hash := fnv.New64a()
	for _, address := range []netip.Addr{
		flow.TupleOrig.IP.SourceAddress,
		flow.TupleOrig.IP.DestinationAddress,
		flow.TupleReply.IP.SourceAddress,
		flow.TupleReply.IP.DestinationAddress,
	} {
		_, _ = hash.Write(address.AsSlice())
	}
	var values [8]byte
	for _, port := range []uint16{
		flow.TupleOrig.Proto.SourcePort,
		flow.TupleOrig.Proto.DestinationPort,
		flow.TupleReply.Proto.SourcePort,
		flow.TupleReply.Proto.DestinationPort,
	} {
		binary.BigEndian.PutUint16(values[:2], port)
		_, _ = hash.Write(values[:2])
	}
	return hash.Sum64()%denominator == 0
}

func tcpState(value uint8) string {
	states := []string{"none", "syn_sent", "syn_recv", "established", "fin_wait", "close_wait", "last_ack", "time_wait", "close", "syn_sent2"}
	if int(value) >= len(states) {
		return fmt.Sprintf("unknown_%d", value)
	}
	return states[value]
}