package media

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"time"

	"git.ipao.vip/rogee/go-sip/internal/audio"
	"github.com/pion/rtp/v2"
	"github.com/zaf/g711"
)

const (
	maxRTPPacketBytes = 2048
	pcm16FrameSamples = 320 // 20 ms at 16 kHz, the canonical AI frame size
)

// Format is the wire codec used by an Asterisk ExternalMedia channel. The
// CallFlow always consumes and produces canonical mono PCM16/16 kHz.
type Format string

const (
	FormatSLIN16 Format = "slin16"
	FormatALAW   Format = "alaw"
)

// RTPStream is a bounded UDP/RTP adapter for an Asterisk ExternalMedia channel.
// Pion owns RTP framing; this package only binds the project media contract and
// exposes signed-linear PCM16 payloads to the AI runtime.
type RTPStream struct {
	conn            *net.UDPConn
	peer            *net.UDPAddr
	payloadType     uint8
	sequence        uint16
	timestamp       uint32
	ssrc            uint32
	format          Format
	wireSampleRate  int
	wireFrameBytes  int
	receivedPackets uint64
	receivedBytes   uint64
	sentPackets     uint64
	sentBytes       uint64
}

type RTPStats struct {
	ReceivedPackets uint64
	ReceivedBytes   uint64
	SentPackets     uint64
	SentBytes       uint64
}

func ListenRTP(address string, payloadType uint8) (*RTPStream, error) {
	return ListenRTPWithFormat(address, payloadType, FormatSLIN16, 16000)
}

func ValidateFormat(format Format, payloadType uint8, sampleRate int) error {
	return validateFormat(format, payloadType, sampleRate)
}

func validateFormat(format Format, payloadType uint8, sampleRate int) error {
	switch format {
	case FormatSLIN16:
		if sampleRate != 16000 || payloadType < 96 || payloadType > 127 {
			return fmt.Errorf("invalid slin16 RTP profile: sample_rate=%d payload_type=%d", sampleRate, payloadType)
		}
	case FormatALAW:
		if sampleRate != 8000 || payloadType != 8 {
			return fmt.Errorf("invalid alaw RTP profile: sample_rate=%d payload_type=%d", sampleRate, payloadType)
		}
	default:
		return fmt.Errorf("unsupported RTP format %q", format)
	}
	return nil
}

func ListenRTPWithFormat(address string, payloadType uint8, format Format, sampleRate int) (*RTPStream, error) {
	if err := validateFormat(format, payloadType, sampleRate); err != nil {
		return nil, err
	}
	addr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		return nil, err
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return nil, err
	}
	var ssrcBytes [4]byte
	if _, err := rand.Read(ssrcBytes[:]); err != nil {
		_ = conn.Close()
		return nil, err
	}
	return &RTPStream{
		conn:           conn,
		payloadType:    payloadType,
		sequence:       1,
		ssrc:           binary.BigEndian.Uint32(ssrcBytes[:]),
		format:         format,
		wireSampleRate: sampleRate,
		wireFrameBytes: sampleRate / 50 * 2,
	}, nil
}

func (s *RTPStream) LocalAddr() net.Addr { return s.conn.LocalAddr() }

func (s *RTPStream) Close() error {
	if s == nil || s.conn == nil {
		return nil
	}
	return s.conn.Close()
}

// SetPeer configures the Asterisk ExternalMedia RTP destination before the
// first inbound packet arrives. A received packet may still refine it to the
// actual source address used by the connected RTP socket.
func (s *RTPStream) SetPeer(address string) error {
	if s == nil || s.conn == nil {
		return errors.New("RTP stream is closed")
	}
	peer, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		return err
	}
	s.peer = peer
	return nil
}

// ReadPayload reads one RTP packet, validates version/payload type and remembers
// the Asterisk peer for subsequent outbound audio.
func (s *RTPStream) ReadPayload(ctx context.Context) ([]byte, error) {
	if s == nil || s.conn == nil {
		return nil, errors.New("RTP stream is closed")
	}
	buf := make([]byte, maxRTPPacketBytes)
	for {
		deadline := time.Now().Add(250 * time.Millisecond)
		if err := s.conn.SetReadDeadline(deadline); err != nil {
			return nil, err
		}
		n, peer, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Timeout() {
				select {
				case <-ctx.Done():
					return nil, ctx.Err()
				default:
					continue
				}
			}
			return nil, err
		}
		var packet rtp.Packet
		if err := packet.Unmarshal(buf[:n]); err != nil {
			continue
		}
		if packet.Version != 2 || packet.PayloadType != s.payloadType {
			continue
		}
		s.peer = peer
		s.receivedPackets++
		s.receivedBytes += uint64(len(packet.Payload))
		payload := append([]byte(nil), packet.Payload...)
		if s.format == FormatALAW {
			payload = audio.ResamplePCM16(g711.DecodeAlaw(payload), s.wireSampleRate, 16000)
		}
		return payload, nil
	}
}

func (s *RTPStream) Stats() RTPStats {
	if s == nil {
		return RTPStats{}
	}
	return RTPStats{ReceivedPackets: s.receivedPackets, ReceivedBytes: s.receivedBytes, SentPackets: s.sentPackets, SentBytes: s.sentBytes}
}

// SendPCM16 sends signed-linear mono PCM16 in 20 ms RTP frames. The peer may
// be configured from Asterisk's UNICASTRTP_LOCAL_* variables before capture.
func (s *RTPStream) SendPCM16(ctx context.Context, pcm []byte, sampleRate int) error {
	if s == nil || s.conn == nil {
		return errors.New("RTP stream is closed")
	}
	if s.peer == nil {
		return errors.New("RTP peer is unknown; configure the ExternalMedia peer first")
	}
	if sampleRate != 16000 {
		return fmt.Errorf("unsupported PCM sample rate %d", sampleRate)
	}
	wirePCM := audio.ResamplePCM16(pcm, sampleRate, s.wireSampleRate)
	for offset := 0; offset < len(wirePCM); {
		end := offset + s.wireFrameBytes
		if s.format == FormatALAW {
			end = offset + s.wireSampleRate/50*2
		}
		if end > len(wirePCM) {
			end = len(wirePCM)
		}
		if end-offset < 2 {
			break
		}
		frame := wirePCM[offset:end]
		payload := frame
		if s.format == FormatALAW {
			payload = g711.EncodeAlaw(frame)
		}
		packet := &rtp.Packet{
			Header: rtp.Header{
				Version:        2,
				PayloadType:    s.payloadType,
				SequenceNumber: s.sequence,
				Timestamp:      s.timestamp,
				SSRC:           s.ssrc,
			},
			Payload: append([]byte(nil), payload...),
		}
		encoded, err := packet.Marshal()
		if err != nil {
			return err
		}
		if err := s.conn.SetWriteDeadline(time.Now().Add(1 * time.Second)); err != nil {
			return err
		}
		if _, err := s.conn.WriteToUDP(encoded, s.peer); err != nil {
			return err
		}
		s.sequence++
		s.timestamp += uint32(len(frame) / 2)
		s.sentPackets++
		s.sentBytes += uint64(len(payload))
		offset = end
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(20 * time.Millisecond):
		}
	}
	return nil
}