package com.usky.cdi.service.sip.rtp;

import com.usky.cdi.service.sip.model.VideoStreamSession;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.OutputStream;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.SocketException;
import java.util.function.Consumer;

@Slf4j
public class RtpVideoReceiver implements Runnable {

    private final VideoStreamSession session;
    private final int localPort;
    private final boolean isPsStream;
    private final Consumer<byte[]> psConsumer;
    private final H264Depacketizer h264Depacketizer;
    private final OutputStream h264OutputStream;

    private volatile boolean running = true;
    private DatagramSocket socket;

    public RtpVideoReceiver(VideoStreamSession session, int localPort, boolean isPsStream,
                            OutputStream h264OutputStream, Consumer<byte[]> psConsumer) {
        this.session = session;
        this.localPort = localPort;
        this.isPsStream = isPsStream;
        this.h264OutputStream = h264OutputStream;
        this.psConsumer = psConsumer;
        this.h264Depacketizer = isPsStream ? null : new H264Depacketizer();
    }

    public void start() throws SocketException {
        socket = new DatagramSocket(localPort);
        socket.setSoTimeout(5000);
        Thread thread = new Thread(this, "rtp-receiver-" + session.getSessionId());
        thread.setDaemon(true);
        thread.start();
        session.setState(VideoStreamSession.SessionState.RECEIVING);
        log.info("RTP 接收已启动: session={}, port={}, codec={}",
                session.getSessionId(), localPort, session.getCodec());
    }

    public void stop() {
        running = false;
        if (socket != null && !socket.isClosed()) {
            socket.close();
        }
        session.setState(VideoStreamSession.SessionState.STOPPED);
        log.info("RTP 接收已停止: session={}", session.getSessionId());
    }

    @Override
    public void run() {
        byte[] buffer = new byte[65535];
        long lastLogPackets = 0;
        while (running) {
            try {
                DatagramPacket datagram = new DatagramPacket(buffer, buffer.length);
                socket.receive(datagram);
                long packets = session.getReceivedPackets().incrementAndGet();
                session.getReceivedBytes().addAndGet(datagram.getLength());

                RtpPacket rtp = RtpPacket.parse(datagram.getData(), datagram.getLength());
                if (rtp == null) {
                    continue;
                }
                if (session.getSsrc() == null) {
                    session.setSsrc(String.format("%08X", rtp.getSsrc()));
                }

                if (isPsStream) {
                    handlePsPayload(rtp);
                } else if (h264OutputStream != null) {
                    h264Depacketizer.depacketize(rtp, h264OutputStream);
                    if (rtp.isMarker()) {
                        h264OutputStream.flush();
                        session.getReceivedFrames().incrementAndGet();
                    }
                }

                if (packets - lastLogPackets >= 100) {
                    lastLogPackets = packets;
                    log.info("RTP 接收统计: callId={}, packets={}, bytes={}",
                            session.getCallId(), packets, session.getReceivedBytes().get());
                }
            } catch (java.net.SocketTimeoutException e) {
                // 超时继续等待
            } catch (IOException e) {
                if (running) {
                    log.warn("RTP 接收异常: session={}, error={}", session.getSessionId(), e.getMessage());
                }
            }
        }
    }

    private void handlePsPayload(RtpPacket rtp) {
        if (psConsumer == null || rtp.getPayloadLength() <= 0) {
            return;
        }
        byte[] chunk = new byte[rtp.getPayloadLength()];
        System.arraycopy(rtp.getPayload(), rtp.getPayloadOffset(), chunk, 0, chunk.length);
        psConsumer.accept(chunk);
    }
}
