package main

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"os"
	"os/exec"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/interceptor"
	"github.com/pion/rtp"
	"github.com/pion/stun/v4"
	"github.com/pion/webrtc/v4"
	"github.com/pion/webrtc/v4/pkg/media"
	"golang.org/x/sys/unix"
	"gopkg.in/hraban/opus.v2"
)

const (
	rate       = 48000
	channels   = 2
	frameSamps = 960 // 20 ms per channel
	opusFmtp   = "minptime=10;useinbandfec=1;stereo=1;sprop-stereo=1;maxaveragebitrate=128000"
)

type Call struct {
	pc     *webrtc.PeerConnection
	track  *webrtc.TrackLocalStaticSample
	ctx    context.Context
	cancel context.CancelFunc
	once   sync.Once

	started  time.Time
	pcState  atomic.Value // string
	iceState atomic.Value // string
	route    atomic.Value // string: selected candidate pair

	outPeak atomic.Uint32 // mix-minus to guest, peak*1e4 since last status
	inPeak  atomic.Uint32 // guest to show
	jb      *jitterBuffer
	jbMu    sync.Mutex
	sendOn  atomic.Bool
	playOn  atomic.Bool
}

var (
	pubIPMu    sync.Mutex
	pubIPCache string
	pubIPAt    time.Time
)

// publicIP returns the rig's public IPv4: -public-ip if set, else a STUN query (works no
// matter how DNS / Cloudflare proxying is set up), else the A record of -public-host.
func publicIP() (string, error) {
	if *publicIPFlag != "" {
		return *publicIPFlag, nil
	}
	pubIPMu.Lock()
	defer pubIPMu.Unlock()
	if pubIPCache != "" && time.Since(pubIPAt) < 10*time.Minute {
		return pubIPCache, nil
	}
	for _, srv := range []string{"stun.l.google.com:19302", "stun.cloudflare.com:3478"} {
		if ip, err := stunIP(srv); err == nil {
			pubIPCache, pubIPAt = ip, time.Now()
			return ip, nil
		} else {
			log.Printf("stun %s: %v", srv, err)
		}
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	ips, err := net.DefaultResolver.LookupIP(ctx, "ip4", *publicHost)
	if err != nil || len(ips) == 0 {
		return "", fmt.Errorf("no public IP (stun failed; resolve %s: %v)", *publicHost, err)
	}
	return ips[0].String(), nil
}

func stunIP(server string) (string, error) {
	conn, err := net.DialTimeout("udp4", server, 2*time.Second)
	if err != nil {
		return "", err
	}
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(2 * time.Second))
	msg := stun.MustBuild(stun.TransactionID, stun.BindingRequest)
	if _, err := conn.Write(msg.Raw); err != nil {
		return "", err
	}
	buf := make([]byte, 1500)
	n, err := conn.Read(buf)
	if err != nil {
		return "", err
	}
	res := &stun.Message{Raw: buf[:n]}
	if err := res.Decode(); err != nil {
		return "", err
	}
	var xor stun.XORMappedAddress
	if err := xor.GetFrom(res); err != nil {
		return "", err
	}
	if xor.IP.To4() == nil {
		return "", errors.New("not ipv4")
	}
	return xor.IP.String(), nil
}

func newAPI() (*webrtc.API, error) {
	m := &webrtc.MediaEngine{}
	if err := m.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus, ClockRate: rate, Channels: 2, SDPFmtpLine: opusFmtp},
		PayloadType:        111,
	}, webrtc.RTPCodecTypeAudio); err != nil {
		return nil, err
	}
	ir := &interceptor.Registry{}
	if err := webrtc.RegisterDefaultInterceptors(m, ir); err != nil {
		return nil, err
	}

	_, lan, err := net.ParseCIDR(*lanCIDR)
	if err != nil {
		return nil, err
	}
	se := webrtc.SettingEngine{}
	se.SetICEUDPMux(hub.mux)
	se.SetNetworkTypes([]webrtc.NetworkType{webrtc.NetworkTypeUDP4})
	se.SetIPFilter(func(ip net.IP) bool { return lan.Contains(ip) })
	se.SetICETimeouts(8*time.Second, 25*time.Second, 2*time.Second)
	if pip, err := publicIP(); err == nil {
		// Keep the LAN host candidate and ADD one with the public IP on the same port.
		// Needs UDP <udp> forwarded to this machine on the router.
		if err := se.SetICEAddressRewriteRules(webrtc.ICEAddressRewriteRule{
			External: []string{pip}, AsCandidateType: webrtc.ICECandidateTypeHost, Mode: webrtc.ICEAddressRewriteAppend,
		}); err != nil {
			return nil, err
		}
	} else {
		log.Printf("warning: %v - offering LAN candidate only", err)
	}
	return webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(ir), webrtc.WithSettingEngine(se)), nil
}

func newCall(offerSDP string) (*Call, string, error) {
	api, err := newAPI()
	if err != nil {
		return nil, "", err
	}
	pc, err := api.NewPeerConnection(webrtc.Configuration{})
	if err != nil {
		return nil, "", err
	}
	ctx, cancel := context.WithCancel(context.Background())
	c := &Call{pc: pc, ctx: ctx, cancel: cancel, started: time.Now()}
	c.pcState.Store("new")
	c.iceState.Store("new")
	c.route.Store("")

	fail := func(err error) (*Call, string, error) { c.Close(); return nil, "", err }

	if err := pc.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: offerSDP}); err != nil {
		return fail(err)
	}
	track, err := webrtc.NewTrackLocalStaticSample(
		webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus, ClockRate: rate, Channels: 2, SDPFmtpLine: opusFmtp},
		"audio", "mmo-mixminus")
	if err != nil {
		return fail(err)
	}
	c.track = track
	sender, err := pc.AddTrack(track)
	if err != nil {
		return fail(err)
	}
	go func() { // drain RTCP so interceptors work
		buf := make([]byte, 1500)
		for {
			if _, _, err := sender.Read(buf); err != nil {
				return
			}
		}
	}()

	pc.OnICEConnectionStateChange(func(s webrtc.ICEConnectionState) {
		c.iceState.Store(s.String())
		log.Printf("ice: %s", s)
		if s == webrtc.ICEConnectionStateConnected {
			c.updateRoute()
		}
	})
	pc.OnConnectionStateChange(func(s webrtc.PeerConnectionState) {
		c.pcState.Store(s.String())
		log.Printf("peer connection: %s", s)
		switch s {
		case webrtc.PeerConnectionStateConnected:
			c.updateRoute()
			c.startSend()
		case webrtc.PeerConnectionStateFailed, webrtc.PeerConnectionStateClosed:
			hub.mu.Lock()
			mine := hub.call == c
			hub.mu.Unlock()
			if mine {
				go hub.endCall()
			}
		}
		go hub.pushState()
	})
	pc.OnTrack(func(t *webrtc.TrackRemote, _ *webrtc.RTPReceiver) {
		if t.Kind() != webrtc.RTPCodecTypeAudio {
			return
		}
		log.Printf("guest track: %s %dHz ch=%d", t.Codec().MimeType, t.Codec().ClockRate, t.Codec().Channels)
		c.startPlay(t)
	})

	answer, err := pc.CreateAnswer(nil)
	if err != nil {
		return fail(err)
	}
	gather := webrtc.GatheringCompletePromise(pc)
	if err := pc.SetLocalDescription(answer); err != nil {
		return fail(err)
	}
	select {
	case <-gather:
	case <-time.After(3 * time.Second):
	}
	return c, pc.LocalDescription().SDP, nil
}

func (c *Call) updateRoute() {
	for _, tr := range c.pc.GetTransceivers() {
		if s := tr.Sender(); s != nil && s.Transport() != nil {
			ice := s.Transport().ICETransport()
			if ice == nil {
				continue
			}
			if pair, err := ice.GetSelectedCandidatePair(); err == nil && pair != nil {
				c.route.Store(fmt.Sprintf("%s:%d <-> %s:%d (%s)", pair.Local.Address, pair.Local.Port, pair.Remote.Address, pair.Remote.Port, pair.Remote.Typ))
				return
			}
		}
	}
}

func (c *Call) Close() {
	c.once.Do(func() {
		c.cancel()
		c.pc.Close()
		log.Printf("call closed after %s", time.Since(c.started).Round(time.Second))
	})
}

func (c *Call) status() map[string]any {
	st := map[string]any{
		"pc":      c.pcState.Load(),
		"ice":     c.iceState.Load(),
		"route":   c.route.Load(),
		"secs":    int(time.Since(c.started).Seconds()),
		"sending": c.sendOn.Load(),
		"playing": c.playOn.Load(),
		"outPeak": float64(c.outPeak.Swap(0)) / 1e4,
		"inPeak":  float64(c.inPeak.Swap(0)) / 1e4,
	}
	c.jbMu.Lock()
	if c.jb != nil {
		st["jb"] = c.jb.stats()
	}
	c.jbMu.Unlock()
	return st
}

func storeMax(a *atomic.Uint32, v uint32) {
	for {
		o := a.Load()
		if v <= o || a.CompareAndSwap(o, v) {
			return
		}
	}
}

func peakOf(pcm []int16) uint32 {
	var p int32
	for _, s := range pcm {
		v := int32(s)
		if v < 0 {
			v = -v
		}
		if v > p {
			p = v
		}
	}
	return uint32(float64(p) / 32768 * 1e4)
}

// ---------------------------------------------------------------- mix-minus -> guest

func (c *Call) startSend() {
	if c.sendOn.Swap(true) {
		return
	}
	go func() {
		defer c.sendOn.Store(false)
		for c.ctx.Err() == nil {
			if err := c.sendLoop(); err != nil && c.ctx.Err() == nil {
				log.Printf("send: %v - retrying in 1s", err)
				select {
				case <-c.ctx.Done():
				case <-time.After(time.Second):
				}
			}
		}
	}()
}

func (c *Call) sendLoop() error {
	cmd := exec.CommandContext(c.ctx, "parec",
		"-d", *sourceName, "--raw", "--format=s16le", fmt.Sprintf("--rate=%d", rate), fmt.Sprintf("--channels=%d", channels),
		"--latency-msec=20", "--property=node.dont-fallback=true", "--property=node.dont-reconnect=true", "--property=node.dont-move=true", "--client-name=guestlink", "--stream-name=mix-minus to guest")
	out, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		return err
	}
	defer cmd.Wait()
	defer cmd.Process.Kill()

	enc, err := opus.NewEncoder(rate, channels, opus.AppAudio)
	if err != nil {
		return err
	}
	enc.SetBitrate(128000)
	enc.SetInBandFEC(true)
	enc.SetPacketLossPerc(5)
	enc.SetComplexity(9)

	raw := make([]byte, frameSamps*channels*2)
	pcm := make([]int16, frameSamps*channels)
	pkt := make([]byte, 4000)
	for {
		if _, err := io.ReadFull(out, raw); err != nil {
			if c.ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("parec read: %w", err)
		}
		for i := range pcm {
			pcm[i] = int16(binary.LittleEndian.Uint16(raw[i*2:]))
		}
		storeMax(&c.outPeak, peakOf(pcm))
		n, err := enc.Encode(pcm, pkt)
		if err != nil {
			return fmt.Errorf("opus encode: %w", err)
		}
		data := make([]byte, n)
		copy(data, pkt[:n])
		if err := c.track.WriteSample(media.Sample{Data: data, Duration: 20 * time.Millisecond}); err != nil && !errors.Is(err, io.ErrClosedPipe) {
			return fmt.Errorf("write sample: %w", err)
		}
	}
}

// ---------------------------------------------------------------- guest -> show

func (c *Call) startPlay(t *webrtc.TrackRemote) {
	dec, err := opus.NewDecoder(rate, channels)
	if err != nil {
		log.Printf("decoder: %v", err)
		return
	}
	jb := newJitterBuffer(dec)
	c.jbMu.Lock()
	c.jb = jb
	c.jbMu.Unlock()

	go func() { // network -> jitter buffer
		for {
			p, _, err := t.ReadRTP()
			if err != nil {
				return
			}
			jb.push(p)
		}
	}()
	go func() { // jitter buffer -> pacat, paced by the MOTU clock via a tiny pipe
		c.playOn.Store(true)
		defer c.playOn.Store(false)
		for c.ctx.Err() == nil {
			if err := c.playLoop(jb); err != nil && c.ctx.Err() == nil {
				log.Printf("play: %v - retrying in 1s", err)
				select {
				case <-c.ctx.Done():
				case <-time.After(time.Second):
				}
			}
		}
	}()
}

func (c *Call) playLoop(jb *jitterBuffer) error {
	if *testPlayoutFile != "" {
		return c.playToFile(jb)
	}
	r, w, err := os.Pipe()
	if err != nil {
		return err
	}
	// Small pipe => writes block on the sink's consumption => the MOTU clock paces us,
	// and the jitter buffer depth tells us about clock drift.
	unix.FcntlInt(w.Fd(), unix.F_SETPIPE_SZ, 4096)
	cmd := exec.CommandContext(c.ctx, "pacat", "--playback",
		"-d", *sinkName, "--raw", "--format=s16le", fmt.Sprintf("--rate=%d", rate), fmt.Sprintf("--channels=%d", channels),
		"--latency-msec=30", "--property=node.dont-fallback=true", "--property=node.dont-reconnect=true", "--property=node.dont-move=true", "--client-name=guestlink", "--stream-name=guest voice")
	cmd.Stdin = r
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		r.Close()
		w.Close()
		return err
	}
	r.Close()
	defer cmd.Wait()
	defer cmd.Process.Kill()
	defer w.Close()

	buf := make([]byte, 0, 5760*channels*2)
	for c.ctx.Err() == nil {
		pcm := jb.next()
		storeMax(&c.inPeak, peakOf(pcm))
		buf = buf[:0]
		for _, s := range pcm {
			buf = binary.LittleEndian.AppendUint16(buf, uint16(s))
		}
		if _, err := w.Write(buf); err != nil {
			if c.ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("pacat write: %w", err)
		}
	}
	return nil
}

// playToFile is a test path: playout to a raw PCM file paced by a 20 ms ticker instead of
// the sound card, so the receive chain can be verified without touching PipeWire.
func (c *Call) playToFile(jb *jitterBuffer) error {
	f, err := os.Create(*testPlayoutFile)
	if err != nil {
		return err
	}
	defer f.Close()
	t := time.NewTicker(20 * time.Millisecond)
	defer t.Stop()
	buf := make([]byte, 0, 5760*channels*2)
	for {
		select {
		case <-c.ctx.Done():
			return nil
		case <-t.C:
		}
		pcm := jb.next()
		storeMax(&c.inPeak, peakOf(pcm))
		buf = buf[:0]
		for _, s := range pcm {
			buf = binary.LittleEndian.AppendUint16(buf, uint16(s))
		}
		f.Write(buf)
	}
}

// ---------------------------------------------------------------- jitter buffer

// jitterBuffer reorders Opus packets, conceals losses (FEC, then PLC), adapts its depth to
// network jitter, and corrects clock drift between the guest's sound card and the MOTU by
// discarding quiet frames when it runs too deep.
type jitterBuffer struct {
	mu       sync.Mutex
	dec      *opus.Decoder
	pkts     map[uint16][]byte
	nextSeq  uint16
	started  bool
	target   int // frames to hold before starting / after an underrun
	misses   int
	lastUnd  time.Time
	pcm      []int16
	lost     int
	fec      int
	plc      int
	dropped  int
	underrun int
	late     int
}

func newJitterBuffer(dec *opus.Decoder) *jitterBuffer {
	return &jitterBuffer{dec: dec, pkts: map[uint16][]byte{}, target: 3, pcm: make([]int16, 5760*channels), lastUnd: time.Now()}
}

func seqLess(a, b uint16) bool { return int16(a-b) < 0 }

func (j *jitterBuffer) push(p *rtp.Packet) {
	j.mu.Lock()
	defer j.mu.Unlock()
	if len(p.Payload) == 0 {
		return
	}
	if j.started {
		if seqLess(p.SequenceNumber, j.nextSeq) {
			j.late++
			return
		}
		if d := uint16(p.SequenceNumber - j.nextSeq); d > 500 { // stream restarted
			j.pkts = map[uint16][]byte{}
			j.started = false
		}
	}
	if len(j.pkts) > 100 { // hard cap - should never happen
		j.pkts = map[uint16][]byte{}
		j.started = false
	}
	j.pkts[p.SequenceNumber] = append([]byte(nil), p.Payload...)
}

func (j *jitterBuffer) oldest() uint16 {
	first := true
	var m uint16
	for s := range j.pkts {
		if first || seqLess(s, m) {
			m, first = s, false
		}
	}
	return m
}

func (j *jitterBuffer) silence() []int16 { return make([]int16, frameSamps*channels) }

func (j *jitterBuffer) decode(p []byte) []int16 {
	n, err := j.dec.Decode(p, j.pcm)
	if err != nil || n <= 0 {
		return j.conceal()
	}
	return append([]int16(nil), j.pcm[:n*channels]...)
}

func (j *jitterBuffer) conceal() []int16 {
	out := make([]int16, frameSamps*channels)
	if err := j.dec.DecodePLC(out); err != nil {
		return j.silence()
	}
	j.plc++
	return out
}

func rms(pcm []int16) float64 {
	if len(pcm) == 0 {
		return 0
	}
	var s float64
	for _, v := range pcm {
		f := float64(v) / 32768
		s += f * f
	}
	return math.Sqrt(s / float64(len(pcm)))
}

// next returns the next block of interleaved PCM to play (normally 20 ms).
func (j *jitterBuffer) next() []int16 {
	j.mu.Lock()
	defer j.mu.Unlock()

	// Relax the target after a quiet minute.
	if j.target > 3 && time.Since(j.lastUnd) > 60*time.Second {
		j.target--
		j.lastUnd = time.Now()
	}

	if !j.started {
		if len(j.pkts) < j.target {
			return j.silence()
		}
		j.nextSeq = j.oldest()
		j.started = true
		j.misses = 0
	}

	for {
		p, ok := j.pkts[j.nextSeq]
		if ok {
			delete(j.pkts, j.nextSeq)
			j.nextSeq++
			j.misses = 0
			pcm := j.decode(p)
			depth := len(j.pkts)
			// Drift / jitter-spike correction: too deep -> throw away quiet frames,
			// or any frame if far too deep.
			if (depth > j.target+3 && rms(pcm) < 0.003) || depth > j.target+12 {
				j.dropped++
				continue
			}
			return pcm
		}
		if len(j.pkts) == 0 {
			// Underrun: nothing to play.
			j.misses++
			if j.misses == 1 {
				j.underrun++
				if j.target < 10 {
					j.target++
				}
				j.lastUnd = time.Now()
			}
			if j.misses > 5 {
				j.started = false // rebuffer to target, play silence meanwhile
				return j.silence()
			}
			j.nextSeq++
			return j.conceal()
		}
		// Gap with later packets present: the packet is lost.
		j.lost++
		if fp, ok := j.pkts[j.nextSeq+1]; ok {
			out := make([]int16, frameSamps*channels)
			if err := j.dec.DecodeFEC(fp, out); err == nil {
				j.fec++
				j.nextSeq++
				return out
			}
		}
		j.nextSeq++
		return j.conceal()
	}
}

func (j *jitterBuffer) stats() map[string]any {
	j.mu.Lock()
	defer j.mu.Unlock()
	return map[string]any{
		"depthMs": len(j.pkts) * 20, "targetMs": j.target * 20,
		"lost": j.lost, "fec": j.fec, "plc": j.plc, "dropped": j.dropped, "underruns": j.underrun, "late": j.late,
	}
}