guzijian / wvp-GB28181-pro · Files

package com.genersoft.iot.vmp.jtt1078.publisher;

import com.genersoft.iot.vmp.jtt1078.rtmp.RtmpChunkWriter;
import com.genersoft.iot.vmp.jtt1078.rtmp.RtmpConnection;
import com.genersoft.iot.vmp.jtt1078.util.Configs;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * 原生Java RTMP推流器
 *
 * 实现了IStreamPublisher接口，使用原生Java实现RTMP协议推送
 *
 * 核心功能：
 * 1. RTMP握手和连接
 * 2. FLV数据到RTMP消息的转换
 * 3. 断流检测和自动重连
 * 4. 码流参数变化时无缝切换（发送新的SPS/PPS）
 */
public class NativeRtmpPublisher implements IStreamPublisher, Runnable {

    static Logger logger = LoggerFactory.getLogger(NativeRtmpPublisher.class);

    private String tag;
    private volatile boolean connected = false;
    private volatile boolean running = false;
    private static final String PUBLISHER_TYPE = "native";

    // RTMP连接配置
    private String host;
    private int port;
    private String app;
    private String tcUrl;
    private String streamName;

    // RTMP组件
    private RtmpConnection connection;
    private RtmpChunkWriter chunkWriter;

    // 重连配置
    private static final int MAX_RECONNECT_ATTEMPTS = 5;
    private static final long INITIAL_RECONNECT_DELAY_MS = 1000;
    private static final long MAX_RECONNECT_DELAY_MS = 16000;

    private final AtomicInteger reconnectAttempts = new AtomicInteger(0);
    private long lastReconnectTime = 0;

    // 心跳检测
    private static final long HEARTBEAT_INTERVAL_MS = 5000;
    private long lastHeartbeatTime = 0;

    // 线程
    private Thread heartbeatThread;
    private volatile Thread currentThread;

    // SPS/PPS缓存（用于无缝切换）
    private byte[] cachedSPS = null;
    private byte[] cachedPPS = null;
    private volatile boolean needSendNewSequenceHeader = false;

    public NativeRtmpPublisher(String tag) {
        this.tag = tag;
        loadConfig();
    }

    /**
     * 从配置加载RTMP连接参数
     *
     * 优先级：
     * 1. rtmp.native.tcUrl (完整TCUrl配置)
     * 2. rtmp.native.host/port/app (分离配置)
     * 3. rtmp.url (兼容旧配置)
     * 4. 默认值
     */
    private void loadConfig() {
        // 优先使用原生RTMP配置
        String configHost = Configs.get("rtmp.native.host");
        String configPort = Configs.get("rtmp.native.port");
        String configApp = Configs.get("rtmp.native.app");
        String configTcUrl = Configs.get("rtmp.native.tcUrl");

        // 如果有完整的tcUrl配置，直接使用
        if (configTcUrl != null && !configTcUrl.isEmpty()) {
            this.tcUrl = configTcUrl.replace("{TAG}", tag);
            // 从tcUrl解析host和port
            parseTcUrl(configTcUrl);
            logger.info("[{}] 使用配置的tcUrl: {}", tag, this.tcUrl);
            return;
        }

        // 如果有分离配置，使用分离配置
        if (configHost != null && !configHost.isEmpty()) {
            this.host = configHost;
        } else {
            this.host = "127.0.0.1";
        }

        if (configPort != null && !configPort.isEmpty()) {
            this.port = Integer.parseInt(configPort);
        } else {
            this.port = 1935;
        }

        if (configApp != null && !configApp.isEmpty()) {
            this.app = configApp;
        } else {
            this.app = "schedule";
        }

        // streamName = tag
        this.streamName = tag;

        // 从rtmp.url获取sign参数（参考FFmpegPublisher的处理方式）
        String sign = "41db35390ddad33f83944f44b8b75ded";  // 默认sign
        String rtmpUrl = Configs.get("rtmp.url");
        if (rtmpUrl != null && rtmpUrl.contains("sign=")) {
            int signStart = rtmpUrl.indexOf("sign=") + 5;
            int signEnd = rtmpUrl.indexOf("&", signStart);
            if (signEnd == -1) signEnd = rtmpUrl.length();
            String extractedSign = rtmpUrl.substring(signStart, signEnd);
            if (!extractedSign.isEmpty() && !"{sign}".equals(extractedSign)) {
                sign = extractedSign;
            }
        }

        // FFmpeg使用完整URL: rtmp://host:port/app/streamName?sign=xxx，ZLM也能正确解析
        // 将streamName和sign包含在tcUrl中
        this.tcUrl = String.format("rtmp://%s:%d/%s/%s?sign=%s", host, port, app, tag, sign);

        logger.info("[{}] RTMP推流器配置: host={}, port={}, app={}, stream={}, tcUrl={}",
                tag, host, port, app, streamName, tcUrl);
    }

    /**
     * 从tcUrl解析host和port
     */
    private void parseTcUrl(String tcUrl) {
        try {
            // 移除sign参数
            if (tcUrl.contains("?")) {
                tcUrl = tcUrl.substring(0, tcUrl.indexOf("?"));
            }

            // 支持 rtmp://, rtsp://, http:// 等协议
            String prefix = "://";
            int prefixIndex = tcUrl.indexOf(prefix);
            if (prefixIndex == -1) {
                logger.warn("[{}] tcUrl没有有效的协议前缀: {}", tag, tcUrl);
                return;
            }

            String afterPrefix = tcUrl.substring(prefixIndex + prefix.length());

            // 分割host:port和path
            int slashIndex = afterPrefix.indexOf("/");
            String hostPort;
            String path;
            if (slashIndex == -1) {
                hostPort = afterPrefix;
                path = "";
            } else {
                hostPort = afterPrefix.substring(0, slashIndex);
                path = afterPrefix.substring(slashIndex + 1);
            }

            // 解析host:port
            int colonIndex = hostPort.indexOf(":");
            if (colonIndex == -1) {
                this.host = hostPort;
                // 默认端口根据协议
                if (tcUrl.startsWith("rtmp://")) {
                    this.port = 1935;
                } else if (tcUrl.startsWith("rtsp://")) {
                    this.port = 554;
                } else {
                    this.port = 1935;
                }
            } else {
                this.host = hostPort.substring(0, colonIndex);
                this.port = Integer.parseInt(hostPort.substring(colonIndex + 1));
            }

            // app是路径的第一部分
            if (path.length() > 0) {
                this.app = path;
                this.streamName = tag;
                this.tcUrl = String.format("rtmp://%s:%d/%s", host, port, path);
            }

        } catch (Exception e) {
            logger.error("[{}] 解析tcUrl失败: {}", tag, e.getMessage());
        }
    }

    /**
     * 解析RTMP URL（兼容方法）
     */
    private void parseRtmpUrl(String rtmpUrl) {
        try {
            // 移除sign参数
            if (rtmpUrl.contains("?")) {
                rtmpUrl = rtmpUrl.substring(0, rtmpUrl.indexOf("?"));
            }

            // 检测协议前缀
            String prefix = "://";
            int prefixIndex = rtmpUrl.indexOf(prefix);
            if (prefixIndex == -1) {
                logger.warn("[{}] URL没有有效的协议前缀: {}", tag, rtmpUrl);
                return;
            }

            String afterPrefix = rtmpUrl.substring(prefixIndex + prefix.length());

            // 分割host:port和app/stream
            int slashIndex = afterPrefix.indexOf("/");
            if (slashIndex == -1) {
                slashIndex = afterPrefix.length();
            }

            String hostPort = afterPrefix.substring(0, slashIndex);
            String appStream = afterPrefix.substring(slashIndex + 1);

            // 解析host:port
            int colonIndex = hostPort.indexOf(":");
            if (colonIndex == -1) {
                this.host = hostPort;
                this.port = 1935;
            } else {
                this.host = hostPort.substring(0, colonIndex);
                this.port = Integer.parseInt(hostPort.substring(colonIndex + 1));
            }

            // app和stream name
            this.app = appStream;
            this.streamName = tag;
            this.tcUrl = String.format("rtmp://%s:%d/%s", host, port, appStream);

            logger.info("[{}] 解析RTMP URL: host={}, port={}, tcUrl={}, stream={}, app={}",
                    tag, host, port, tcUrl, streamName, app);

        } catch (Exception e) {
            logger.error("[{}] 解析RTMP URL失败: {}, 使用默认值", tag, e.getMessage());
            this.host = "127.0.0.1";
            this.port = 1935;
            this.app = "schedule";
            this.tcUrl = "rtmp://127.0.0.1:1935/schedule";
            this.streamName = tag;
        }
    }

    @Override
    public void start(String tag) {
        this.currentThread = Thread.currentThread();
        this.running = true;

        logger.info("[{}] ====== 启动原生RTMP推流器 ======", tag);
        logger.info("[{}] 目标服务器: {}:{}", tag, host, port);
        logger.info("[{}] 推流地址: {}/{}", tag, tcUrl, streamName);
        logger.info("[{}] 完整推流URL: {}/{}?sign=xxx", tag, tcUrl, streamName);

        // 启动心跳线程
        startHeartbeat();

        // 执行连接
        doConnect();
    }

    @Override
    public void run() {
        start(tag);
    }

    /**
     * 执行连接
     */
    private void doConnect() {
        try {
            // 关闭旧连接
            closeConnection();

            // 创建新连接
            connection = new RtmpConnection(host, port, tcUrl, streamName);

            if (!connection.connect()) {
                logger.error("[{}] RTMP连接失败", tag);
                scheduleReconnect();
                return;
            }

            if (!connection.createStream()) {
                logger.error("[{}] RTMP创建流失败", tag);
                scheduleReconnect();
                return;
            }

            if (!connection.publish()) {
                logger.error("[{}] RTMP发布失败", tag);
                scheduleReconnect();
                return;
            }

            // 【关键修复】等待publish命令被ZLM处理完成
            // jtt1078-video-server项目成功的原因：RtmpHandshakeHandler会等待服务器响应
            // "NetStream.Publish.Start"后才进入STREAMING状态，才开始发送数据
            // 当前项目需要等待一段时间确保ZLM完成publish处理
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }

            // 创建ChunkWriter
            chunkWriter = new RtmpChunkWriter(connection, tag);

            // 如果有缓存的SPS/PPS，标记需要发送
            if (cachedSPS != null && cachedPPS != null) {
                needSendNewSequenceHeader = true;
            }

            connected = true;
            reconnectAttempts.set(0);
            lastHeartbeatTime = System.currentTimeMillis();

            logger.info("[{}] ====== RTMP推流连接成功 ======", tag);

        } catch (Exception e) {
            logger.error("[{}] RTMP连接异常: {}", tag, e.getMessage(), e);
            scheduleReconnect();
        }
    }

    /**
     * 启动心跳检测线程
     */
    private void startHeartbeat() {
        if (heartbeatThread != null && heartbeatThread.isAlive()) {
            return;
        }

        heartbeatThread = new Thread(() -> {
            logger.info("[{}] 心跳检测线程启动", tag);
            while (running && !Thread.currentThread().isInterrupted()) {
                try {
                    Thread.sleep(HEARTBEAT_INTERVAL_MS);

                    if (!running) {
                        break;
                    }

                    // 检查连接状态
                    if (!connected || connection == null || !connection.isConnected()) {
                        logger.warn("[{}] 心跳检测: 连接已断开，尝试重连", tag);
                        scheduleReconnect();
                        break;
                    }

                    // 检查最后心跳时间
                    long now = System.currentTimeMillis();
                    if (now - lastHeartbeatTime > HEARTBEAT_INTERVAL_MS * 3) {
                        logger.warn("[{}] 心跳超时: {}ms未收到数据", tag, now - lastHeartbeatTime);
                        // 连接可能还活着，只是没数据
                    }

                    lastHeartbeatTime = now;

                } catch (InterruptedException e) {
                    logger.info("[{}] 心跳线程被中断", tag);
                    break;
                } catch (Exception e) {
                    logger.error("[{}] 心跳检测异常: {}", tag, e.getMessage());
                }
            }
            logger.info("[{}] 心跳检测线程退出", tag);
        }, "RTMP-Heartbeat-" + tag);
        heartbeatThread.setDaemon(true);
        heartbeatThread.start();
    }

    /**
     * 安排重连
     */
    private void scheduleReconnect() {
        if (!running) {
            return;
        }

        int attempts = reconnectAttempts.incrementAndGet();
        if (attempts > MAX_RECONNECT_ATTEMPTS) {
            logger.error("[{}] 超过最大重连次数({})，停止重连", tag, MAX_RECONNECT_ATTEMPTS);
            logger.info("[{}] 请检查网络和RTMP服务器状态", tag);
            return;
        }

        // 计算退避延迟
        long delay = Math.min(
            INITIAL_RECONNECT_DELAY_MS * (1L << (attempts - 1)),
            MAX_RECONNECT_DELAY_MS
        );

        // 避免过于频繁的重连
        long now = System.currentTimeMillis();
        if (now - lastReconnectTime < delay && attempts > 1) {
            delay = Math.max(delay, now - lastReconnectTime);
        }
        lastReconnectTime = now;

        logger.info("[{}] 计划{}ms后进行第{}次重连 (最多{}次)",
                tag, delay, attempts, MAX_RECONNECT_ATTEMPTS);

        // 使用当前线程执行延迟重连
        final long finalDelay = delay;
        new Thread(() -> {
            try {
                Thread.sleep(finalDelay);
                if (running) {
                    doConnect();
                }
            } catch (InterruptedException e) {
                logger.info("[{}] 重连延迟线程被中断", tag);
            }
        }, "RTMP-Reconnect-" + tag).start();
    }

    /**
     * 关闭连接
     */
    private void closeConnection() {
        connected = false;

        if (connection != null) {
            try {
                connection.close();
            } catch (Exception e) {
                logger.debug("[{}] 关闭连接时出错: {}", tag, e.getMessage());
            }
            connection = null;
        }
    }

    @Override
    public void sendVideoData(byte[] flvData, int timestamp) {
        if (!connected || chunkWriter == null) {
            logger.info("[{}] sendVideoData: not connected or chunkWriter is null, connected={}, chunkWriter={}",
                    tag, connected, chunkWriter);
            return;
        }

        try {
            lastHeartbeatTime = System.currentTimeMillis();
            logger.info("[{}] sendVideoData: flvDataLen={}, timestamp={}", tag, flvData.length, timestamp);
            chunkWriter.sendVideoData(flvData, timestamp);
        } catch (Exception e) {
            logger.error("[{}] 发送视频数据失败: {}", tag, e.getMessage());
            connected = false;
            scheduleReconnect();
        }
    }

    @Override
    public void sendAVCSequenceHeader() {
        if (!connected || chunkWriter == null) {
            logger.warn("[{}] 未连接，无法发送AVC Sequence Header", tag);
            return;
        }

        try {
            // 强制重新发送SPS/PPS
            chunkWriter.forceSendAVCSequenceHeader();
            logger.info("[{}] 已触发发送新的AVC Sequence Header", tag);
        } catch (Exception e) {
            logger.error("[{}] 发送AVC Sequence Header失败: {}", tag, e.getMessage());
        }
    }

    @Override
    public void close() {
        logger.info("[{}] ====== 关闭原生RTMP推流器 ======", tag);

        running = false;
        connected = false;

        // 中断心跳线程
        if (heartbeatThread != null) {
            heartbeatThread.interrupt();
            try {
                heartbeatThread.join(2000);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }

        // 关闭连接
        closeConnection();

        // 中断当前线程
        if (currentThread != null && currentThread.isAlive()) {
            currentThread.interrupt();
        }

        logger.info("[{}] 原生RTMP推流器已关闭", tag);
    }

    @Override
    public boolean isConnected() {
        return connected && connection != null && connection.isConnected();
    }

    @Override
    public String getType() {
        return PUBLISHER_TYPE;
    }

    @Override
    public String getStatus() {
        if (connected) {
            return String.format("原生RTMP推流器, 已连接, streamId=%d, 重连次数=%d",
                    connection != null ? connection.getStreamId() : 0,
                    reconnectAttempts.get());
        } else {
            return String.format("原生RTMP推流器, 未连接, 重连次数=%d/%d",
                    reconnectAttempts.get(), MAX_RECONNECT_ATTEMPTS);
        }
    }

    /**
     * 更新SPS/PPS缓存（当Channel检测到SPS变化时调用）
     */
    public void updateSPSPPS(byte[] sps, byte[] pps) {
        this.cachedSPS = sps;
        this.cachedPPS = pps;
        this.needSendNewSequenceHeader = true;

        if (chunkWriter != null) {
            chunkWriter.updateSPSPPS(sps, pps);
        }
    }

    /**
     * 获取重连次数
     */
    public int getReconnectAttempts() {
        return reconnectAttempts.get();
    }
}