package pipewire

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Capture-jitter telemetry: log a stats line every readStatsInterval reads
// (avg/min/max ReadFull latency) plus an immediate warning whenever a single
// read exceeds readSpikeThreshold. pw-cat produces samples in real time, so
// the expected per-read latency is ~ChunkDurationMs (20ms). Spikes there are
// the most likely on-device cause of audible drops.
const (
	readSpikeThreshold = 40 * time.Millisecond
	readStatsInterval  = 250 // ~5s at 20ms chunks
)

// Source implements the sendspin AudioSource interface via pw-cat --record.
//
// To avoid PipeWire's session manager auto-connecting our recorder to whatever
// the "default" source is (typically the mic on hardware that has one — which
// silently steals capture from a requested `.monitor` target), we start pw-cat
// with node.autoconnect=false and a unique node.name, then look up both nodes
// in pw-dump and pw-link the target's output ports to our input ports
// manually.
type Source struct {
	target     string
	sampleRate int
	channels   int

	mu     sync.Mutex
	cmd    *exec.Cmd
	stdout io.ReadCloser
	buf    []byte

	// dataCh carries raw byte chunks from a background goroutine that
	// drains pw-cat as fast as bytes appear. Read assembles ChunkDurationMs
	// worth of samples from this stream — without this decoupling, Read
	// would block ~21ms per call because PipeWire's default graph quantum
	// (1024 samples at 48kHz) exceeds our 20ms tick budget, pushing the
	// engine ~1.3ms behind real-time per chunk and eroding the 500ms
	// client buffer-ahead until chunks land late and audio drops.
	dataCh    chan []byte
	leftover  []byte
	stopDrain chan struct{}
	drainErr  error
	drainMu   sync.Mutex

	// jitter telemetry — accessed only from Read so no extra locking
	readCount    int
	readSumNanos int64
	readMaxNanos int64
	readMinNanos int64
}

func NewSource(target string, sampleRate, channels int) *Source {
	return &Source{
		target:     target,
		sampleRate: sampleRate,
		channels:   channels,
	}
}

func (s *Source) start() error {
	uid := fmt.Sprintf("sendspin-cap-%d", time.Now().UnixNano())
	uidJSON, _ := json.Marshal(uid)
	// node.latency requests a 20ms quantum (960 samples at 48kHz). Without
	// this, PipeWire defaults to its graph quantum (typically 1024 samples
	// / 21.33ms), which means pw-cat delivers in bursts that don't align
	// with the sendspin engine's 20ms ticker — every chunk ends up 1.3ms
	// over budget and the 500ms client buffer-ahead drifts away.
	chunkSamples := (s.sampleRate * 20) / 1000
	props := fmt.Sprintf(
		`{"node.name":%s,"node.description":"Sendspin capture","node.autoconnect":false,"node.latency":"%d/%d"}`,
		uidJSON, chunkSamples, s.sampleRate,
	)

	args := []string{
		"--record",
		"--raw",
		"--format", "s32",
		"--rate", strconv.Itoa(s.sampleRate),
		"--channels", strconv.Itoa(s.channels),
		"--target", s.target,
		"--properties", props,
		"-",
	}

	log.Printf("pw-cat %s", strings.Join(args, " "))
	cmd := exec.Command("pw-cat", args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("pw-cat stdout: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("pw-cat stderr: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("pw-cat start: %w", err)
	}

	go func(t string) {
		buf, _ := io.ReadAll(stderr)
		if msg := strings.TrimSpace(string(buf)); msg != "" {
			log.Printf("pw-cat[%s] stderr: %s", t, msg)
		}
	}(s.target)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	outIDs, inIDs, linkErr := waitAndMatchPorts(ctx, s.target, uid)
	if linkErr != nil {
		_ = stdout.Close()
		if cmd.Process != nil {
			_ = cmd.Process.Kill()
		}
		_ = cmd.Wait()
		return fmt.Errorf("link target ports: %w", linkErr)
	}

	for i := range outIDs {
		lc := exec.Command("pw-link", strconv.Itoa(outIDs[i]), strconv.Itoa(inIDs[i]))
		if out, err := lc.CombinedOutput(); err != nil {
			log.Printf("pw-link %d→%d: %v: %s", outIDs[i], inIDs[i], err, out)
		}
	}

	s.cmd = cmd
	s.stdout = stdout
	s.buf = make([]byte, 4096)

	// Buffered channel sized to ~1s of audio (50 × 20ms chunks). Plenty
	// of headroom for pw-cat's quantum bursts; if it ever fills, pw-cat
	// blocks on its stdout write — at which point the engine is so far
	// behind that backpressure is the right behavior.
	s.dataCh = make(chan []byte, 50)
	s.stopDrain = make(chan struct{})
	go s.drain(stdout, s.stopDrain)

	// Pre-buffer ~60ms (3 quanta) before letting Read return. Two quanta
	// (~43ms) is the minimum to absorb the 1024-vs-960-sample beat
	// pattern; 60ms gives a small safety margin for scheduler jitter
	// without adding meaningful startup latency.
	preBufferBytes := (60 * s.sampleRate / 1000) * 4 * s.channels
	prebufStart := time.Now()
	have := 0
	var queued [][]byte
	for have < preBufferBytes {
		select {
		case b := <-s.dataCh:
			queued = append(queued, b)
			have += len(b)
		case <-time.After(2 * time.Second):
			_ = stdout.Close()
			if cmd.Process != nil {
				_ = cmd.Process.Kill()
			}
			_ = cmd.Wait()
			return fmt.Errorf("pw-cat pre-buffer timeout (got %d/%d bytes in %s)",
				have, preBufferBytes, time.Since(prebufStart))
		}
	}
	// Re-prime the channel with what we already drained so Read sees
	// the full pre-buffer. Channel capacity (50) > expected queue size
	// (~5 slices), so this never blocks.
	for _, b := range queued {
		s.dataCh <- b
	}
	log.Printf("pw-cat pre-buffered %d bytes (%dms) in %s",
		have, have/(48*4*2), time.Since(prebufStart).Round(time.Millisecond))
	return nil
}

// drain continuously reads from pw-cat's stdout and pushes byte slices
// onto dataCh. Runs until stdout closes or stop is signaled. Each iteration
// allocates a fresh slice so the receiver owns it — pool/reuse here would
// race with Read.
func (s *Source) drain(stdout io.ReadCloser, stop <-chan struct{}) {
	buf := make([]byte, 8192)
	for {
		n, err := stdout.Read(buf)
		if n > 0 {
			chunk := make([]byte, n)
			copy(chunk, buf[:n])
			select {
			case s.dataCh <- chunk:
			case <-stop:
				return
			}
		}
		if err != nil {
			s.drainMu.Lock()
			s.drainErr = err
			s.drainMu.Unlock()
			close(s.dataCh)
			return
		}
		select {
		case <-stop:
			return
		default:
		}
	}
}

func (s *Source) Read(samples []int32) (int, error) {
	s.mu.Lock()
	if s.cmd == nil {
		if err := s.start(); err != nil {
			s.mu.Unlock()
			return 0, err
		}
	}

	need := len(samples) * 4
	if len(s.buf) < need {
		s.buf = make([]byte, need)
	}
	buf := s.buf[:need]
	s.mu.Unlock()

	readStart := time.Now()
	filled := 0

	// Drain any leftover bytes from the previous Read first. pw-cat's
	// quantum (1024 samples) doesn't align with our request (960 samples),
	// so every other call carries 64 samples' worth of bytes forward.
	if len(s.leftover) > 0 {
		copied := copy(buf[filled:], s.leftover)
		filled += copied
		s.leftover = s.leftover[copied:]
	}

	for filled < need {
		chunk, ok := <-s.dataCh
		if !ok {
			s.drainMu.Lock()
			err := s.drainErr
			s.drainMu.Unlock()
			s.mu.Lock()
			s.killUnlocked()
			s.mu.Unlock()
			if err == nil {
				err = io.EOF
			}
			return filled / 4, err
		}
		copied := copy(buf[filled:], chunk)
		filled += copied
		if copied < len(chunk) {
			s.leftover = chunk[copied:]
		}
	}

	s.recordReadTiming(time.Since(readStart), need)

	count := filled / 4
	// pw-cat emits full-range int32; sendspin's AudioSource contract is
	// int32 samples right-justified to 24 bits (encodePCM packs the low
	// three bytes; opus does `int16(s >> 8)`). Without this shift the
	// signal's top byte is discarded and only low-byte dither survives,
	// producing noise that tracks input amplitude.
	for i := range count {
		samples[i] = int32(binary.LittleEndian.Uint32(buf[i*4:i*4+4])) >> 8
	}
	return count, nil
}

// recordReadTiming accumulates per-read latency and emits both spike
// warnings and periodic stats summaries. The expected baseline is roughly
// the chunk duration the engine asks for, since pw-cat produces real-time
// samples and io.ReadFull blocks until that many bytes are available.
func (s *Source) recordReadTiming(dur time.Duration, byteCount int) {
	d := dur.Nanoseconds()
	s.readCount++
	s.readSumNanos += d
	if d > s.readMaxNanos {
		s.readMaxNanos = d
	}
	if s.readMinNanos == 0 || d < s.readMinNanos {
		s.readMinNanos = d
	}

	if dur > readSpikeThreshold {
		log.Printf("pw-cat read spike: %s for %d bytes (threshold %s)",
			dur.Round(time.Microsecond), byteCount, readSpikeThreshold)
	}

	if s.readCount >= readStatsInterval {
		avg := time.Duration(s.readSumNanos / int64(s.readCount))
		log.Printf("pw-cat read stats over %d reads: avg=%s min=%s max=%s",
			s.readCount,
			avg.Round(time.Microsecond),
			time.Duration(s.readMinNanos).Round(time.Microsecond),
			time.Duration(s.readMaxNanos).Round(time.Microsecond),
		)
		s.readCount = 0
		s.readSumNanos = 0
		s.readMaxNanos = 0
		s.readMinNanos = 0
	}
}

// killUnlocked tears down the running pw-cat without holding the mutex.
// Callers must hold s.mu.
func (s *Source) killUnlocked() {
	if s.stopDrain != nil {
		// Signal the drainer to exit. It will also exit on its own when
		// stdout closes, but signaling first avoids a hung send into
		// dataCh if the buffer is full at shutdown.
		close(s.stopDrain)
		s.stopDrain = nil
	}
	if s.stdout != nil {
		s.stdout.Close()
		s.stdout = nil
	}
	if s.cmd != nil && s.cmd.Process != nil {
		s.cmd.Process.Kill()
		s.cmd.Wait()
		s.cmd = nil
	}
	s.leftover = nil
}

func (s *Source) SampleRate() int                    { return s.sampleRate }
func (s *Source) Channels() int                      { return s.channels }
func (s *Source) Metadata() (string, string, string) { return "", "", "" }

func (s *Source) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.killUnlocked()
	return nil
}

type portEntry struct {
	id      int
	channel string
}

type dumpObject struct {
	ID   int             `json:"id"`
	Type string          `json:"type"`
	Info json.RawMessage `json:"info"`
}

type dumpNodeInfo struct {
	Props json.RawMessage `json:"props"`
}

type dumpPortInfo struct {
	Direction string          `json:"direction"`
	Props     json.RawMessage `json:"props"`
}

// waitAndMatchPorts polls pw-dump until both the target node and our pw-cat
// node are present, then pairs target output ports to our input ports.
func waitAndMatchPorts(ctx context.Context, targetName, ourName string) (outIDs, inIDs []int, err error) {
	for {
		select {
		case <-ctx.Done():
			return nil, nil, ctx.Err()
		default:
		}

		raw, dumpErr := exec.Command("pw-dump").Output()
		if dumpErr != nil {
			time.Sleep(100 * time.Millisecond)
			continue
		}
		var objects []dumpObject
		if err := json.Unmarshal(raw, &objects); err != nil {
			time.Sleep(100 * time.Millisecond)
			continue
		}

		nodeIDs := make(map[string]int)
		for _, obj := range objects {
			if obj.Type != "PipeWire:Interface:Node" {
				continue
			}
			var info dumpNodeInfo
			if json.Unmarshal(obj.Info, &info) != nil {
				continue
			}
			var props map[string]any
			if json.Unmarshal(info.Props, &props) != nil {
				continue
			}
			if name, ok := props["node.name"].(string); ok {
				nodeIDs[name] = obj.ID
			}
		}

		targetID, targetOK := nodeIDs[targetName]
		ourID, ourOK := nodeIDs[ourName]
		if !targetOK || !ourOK {
			time.Sleep(100 * time.Millisecond)
			continue
		}

		var targetOuts, ourIns []portEntry
		for _, obj := range objects {
			if obj.Type != "PipeWire:Interface:Port" {
				continue
			}
			var info dumpPortInfo
			if json.Unmarshal(obj.Info, &info) != nil {
				continue
			}
			var props map[string]any
			if json.Unmarshal(info.Props, &props) != nil {
				continue
			}
			nodeIDf, ok := props["node.id"].(float64)
			if !ok {
				continue
			}
			nodeID := int(nodeIDf)
			ch, _ := props["audio.channel"].(string)

			switch {
			case nodeID == targetID && info.Direction == "output":
				targetOuts = append(targetOuts, portEntry{obj.ID, ch})
			case nodeID == ourID && info.Direction == "input":
				ourIns = append(ourIns, portEntry{obj.ID, ch})
			}
		}

		if len(targetOuts) == 0 || len(ourIns) == 0 {
			time.Sleep(100 * time.Millisecond)
			continue
		}

		pairs := matchPorts(targetOuts, ourIns)
		if len(pairs) == 0 {
			return nil, nil, fmt.Errorf("no matching ports between %q and %q", targetName, ourName)
		}
		for _, p := range pairs {
			outIDs = append(outIDs, p[0])
			inIDs = append(inIDs, p[1])
		}
		return outIDs, inIDs, nil
	}
}

// matchPorts pairs output ports to input ports by audio.channel name,
// falling back to positional order when channel names are missing or unmatched.
func matchPorts(outs, ins []portEntry) [][2]int {
	inByChannel := make(map[string]int, len(ins))
	for _, p := range ins {
		if p.channel != "" {
			inByChannel[p.channel] = p.id
		}
	}

	used := make(map[int]bool)
	var result [][2]int
	for _, o := range outs {
		if inID, ok := inByChannel[o.channel]; ok && !used[inID] {
			result = append(result, [2]int{o.id, inID})
			used[inID] = true
		}
	}
	if len(result) > 0 {
		return result
	}

	for i := range min(len(outs), len(ins)) {
		result = append(result, [2]int{outs[i].id, ins[i].id})
	}
	return result
}