#include <iostream>
#include <vector>
#include <cstdint>
#include <cstring>
#include <iomanip>  // 用于十六进制输出格式化
#include <ctime>    // 用于日志中的时间戳
#include <sstream>  // 用于字符串流
#include <netinet/in.h> // 包含 htons 的声明
#include <chrono>
#include <thread>
#include <numeric> // std::iota
#include <malloc.h>

// --- 常量定义 ---
const uint16_t VIDEO_STREAM_ID = 0xE0; // 视频流的 Stream ID 示例
const size_t MAX_RTP_PAYLOAD_SIZE = 1400; // RTP 负载最大尺寸示例
const uint32_t PS_START_CODE_PREFIX = 0x000001BA;
const uint32_t PES_START_CODE_PREFIX = 0x000001;
const uint8_t NAL_START_CODE_4_BYTE[] = {0x00, 0x00, 0x00, 0x01};

const int PS_HDR_LEN = 14; // 示例 PS Header Length
const int SYS_HDR_LEN = 8; // 示例 System Header Length
const int PSM_HDR_LEN = 12; // 示例 PSM Header Length
const int PES_HDR_LEN = 19; // 示例 PES Header Length
const int RTP_HDR_LEN = 12; // 示例 RTP Header Length


// --- 结构体定义 ---
#pragma pack(push, 1) // 确保结构体内部没有填充字节

// 简化的 PS 头部 (示例)
struct PSHeader {
    uint32_t packet_start_code_prefix;
    uint64_t system_clock_reference; // SCR

    PSHeader() : packet_start_code_prefix(PS_START_CODE_PREFIX), system_clock_reference(0) {}
};

// 简化的系统头部 (示例)
struct SystemHeader {
    uint32_t system_header_start_code;
    uint16_t header_length;
    uint8_t rate_bound[3];
    uint8_t audio_bound;
    uint8_t fixed_flag_etc;

    SystemHeader() : system_header_start_code(0x000001BB), header_length(0), audio_bound(0), fixed_flag_etc(0) {
        rate_bound[0] = rate_bound[1] = rate_bound[2] = 0;
    }
};

// 简化的 PSM 头部 (示例)
struct PSMHeader {
    uint32_t program_stream_map_start_code;
    uint16_t psm_length;
    uint8_t program_number[2];
    uint8_t version_current_next_indicator;
    uint8_t section_number;
    uint8_t last_section_number;
    uint8_t program_info_length[2];
    // ... (省略 Program Stream Info 和 ES Info 循环) ...

    PSMHeader() : program_stream_map_start_code(0x000001BC), psm_length(0), version_current_next_indicator(0), section_number(0), last_section_number(0) {
        program_number[0] = program_number[1] = 0;
        program_info_length[0] = program_info_length[1] = 0;
    }
};


// 简化的 PES 头部 (关注 PTS)
struct PESHeader {
    uint32_t packet_start_code_prefix;
    uint8_t stream_id;
    uint16_t pes_packet_length;
    uint8_t pes_scrambling_control_indicator_etc;
    uint8_t pes_header_data_length;
    uint64_t pts_dts_flags_pts;

    PESHeader() : packet_start_code_prefix(PES_START_CODE_PREFIX), stream_id(VIDEO_STREAM_ID), pes_packet_length(0),
                  pes_scrambling_control_indicator_etc(0x80), // PTS_DTS_flags: 0b10 (仅 PTS)
                  pes_header_data_length(5),
                  pts_dts_flags_pts(0) {}
};

// 简化的 RTP 头部
struct RTPHeader {
    uint8_t version_padding_extension_csrc_count; // V, P, X, CC
    uint8_t marker_payload_type;                 // M, PT
    uint16_t sequence_number;
    uint32_t timestamp;
    uint32_t ssrc;

    RTPHeader() : version_padding_extension_csrc_count(0x80), marker_payload_type(96), sequence_number(0), timestamp(0), ssrc(0x12345678) {} // SSRC 示例
};

// Nalu 结构体
struct Nalu {
    int type;
    int length;
    std::vector<uint8_t> packet;

    Nalu() : type(0), length(0) {}
    ~Nalu() {}
};

using NaluType = int;


#pragma pack(pop) // 恢复默认 packing

// --- 全局计数器和变量 ---
static uint64_t pes_pts_counter = 0; // PES PTS 计数器示例
static uint16_t rtp_seq_counter = 0;

// --- 日志记录函数 ---
void Log(const std::string& message) {
    std::time_t now = std::time(nullptr);
    std::tm local_time;
    localtime_r(&now, &local_time);
    char timestamp_str[20];
    std::strftime(timestamp_str, sizeof(timestamp_str), "%Y-%m-%d %H:%M:%S", &local_time);
    std::cout << "[" << timestamp_str << "] " << message << std::endl;
}

// --- 辅助函数：将数据转换为十六进制字符串 ---
std::string ToHex(const uint8_t* data, size_t size) {
    std::stringstream hex_stream;
    hex_stream << std::hex << std::setfill('0');
    for (size_t i = 0; i < size; ++i) {
        hex_stream << std::setw(2) << static_cast<int>(data[i]) << " ";
    }
    return hex_stream.str();
}

// --- 将 VPS, SPS, PPS 处理为一个流 ---
std::vector<uint8_t> vps_data;
std::vector<uint8_t> sps_data;
std::vector<uint8_t> pps_data;

void out_nalu(char *buffer, int size, NaluType naluType) {
    if (naluType == 32) {  // VPS
        vps_data.resize(size);
        memcpy(vps_data.data(), buffer, size);
    } else if (naluType == 33) {  // SPS
        sps_data.resize(size);
        memcpy(sps_data.data(), buffer, size);
    } else if (naluType == 34) {  // PPS
        pps_data.resize(size);
        memcpy(pps_data.data(), buffer, size);
    } else {
        Nalu *nalu = new Nalu;
        bool is_i_frame = (naluType == 19);  // IDR frame
        char *packet = (char *)malloc(is_i_frame ? (size + vps_data.size() + sps_data.size() + pps_data.size()) : size);

        if (is_i_frame) {
            memcpy(packet, vps_data.data(), vps_data.size());
            memcpy(packet + vps_data.size(), sps_data.data(), sps_data.size());
            memcpy(packet + vps_data.size() + sps_data.size(), pps_data.data(), pps_data.size());
            memcpy(packet + vps_data.size() + sps_data.size() + pps_data.size(), buffer, size);
            size += (vps_data.size() + sps_data.size() + pps_data.size());
        } else {
            memcpy(packet, buffer, size);
        }

        nalu->packet = std::vector<uint8_t>(packet, packet + size);
        nalu->length = size;
        nalu->type = naluType;
        // 将 nalu 添加到 nalu_vector 或进行其他处理
        delete[] packet;
    }
}

// --- 头部生成函数 ---
void gb28181_make_ps_header(char *header, long pts) {
    PSHeader ps_header_struct;
    ps_header_struct.system_clock_reference = pts; // 简化 SCR
    memcpy(header, &ps_header_struct, sizeof(PSHeader));
}

void gb28181_make_sys_header(char *header, int rate_bound) {
    SystemHeader sys_header_struct;
    sys_header_struct.header_length = htons(SYS_HDR_LEN - 6); // Length after header_length field
    sys_header_struct.rate_bound[0] = (rate_bound >> 16) & 0xFF;
    sys_header_struct.rate_bound[1] = (rate_bound >> 8) & 0xFF;
    sys_header_struct.rate_bound[2] = rate_bound & 0xFF;

    memcpy(header, &sys_header_struct, sizeof(SystemHeader));
}

void gb28181_make_psm_header(char *header) {
    PSMHeader psm_header_struct;
    psm_header_struct.psm_length = htons(PSM_HDR_LEN - 6); // Length after psm_length field
    memcpy(header, &psm_header_struct, sizeof(PSMHeader));
}

void gb28181_make_pes_header(char *header, int stream_id, int data_len, long pts, long dts) {
    PESHeader pes_header_struct;
    pes_header_struct.stream_id = stream_id;
    pes_header_struct.pes_packet_length = htons(data_len + PES_HDR_LEN - 6); // Length after pes_packet_length field
    pes_header_struct.pts_dts_flags_pts = (static_cast<uint64_t>(pes_pts_counter++) << 3) | 0x02; // 仅设置 PTS 标志
    uint64_t pts_33_to_1 = (pes_pts_counter * 300) % 0x200000000LL; // 假设 300 ticks/ms, 90kHz clock
    uint32_t pts_32_to_2 = pts_33_to_1 & 0xFFFFFFFE0LL;
    uint8_t pts_byte0 = 0x20 | ((pts_32_to_2 >> 30) & 0x07) << 1 | 0x01;
    uint16_t pts_byte1_2 = (pts_32_to_2 >> 15) & 0xFFFF;
    uint16_t pts_byte3_4 = pts_32_to_2 & 0xFFFF;

    pes_header_struct.pts_dts_flags_pts |= (static_cast<uint64_t>(pts_byte0) << 40);
    pes_header_struct.pts_dts_flags_pts |= (static_cast<uint64_t>(pts_byte1_2) << 24);
    pes_header_struct.pts_dts_flags_pts |= (static_cast<uint64_t>(pts_byte3_4) >> 8);

    memcpy(header, &pes_header_struct, sizeof(PESHeader));
}

void gb28181_make_rtp_header(char *header, int seq, long pts, int ssrc, bool marker) {
    RTPHeader rtp_header_struct;
    rtp_header_struct.sequence_number = htons(seq);
    rtp_header_struct.timestamp = htonl(pts);
    rtp_header_struct.ssrc = htonl(ssrc);
    if (marker) {
        rtp_header_struct.marker_payload_type |= (1 << 7); // Set marker bit
    }
    memcpy(header, &rtp_header_struct, sizeof(RTPHeader));
}

// --- 发送网络数据包 ---
void send_network_packet(char *packet, int packet_size) {
    Log("发送网络包, 大小: " + std::to_string(packet_size) + " bytes.");
    // 模拟发送网络包，实际应用中替换为 socket send 操作
    Log("Packet Data (First 64 bytes): \n" + ToHex((uint8_t*)packet, std::min((size_t)64, (size_t)packet_size)));
}

// --- 判断是否为关键帧 ---
bool is_keyframe(NaluType type) {
    // 这里可以根据 NALU 类型判断是否为关键帧
    // 例如，对于 H.264, IDR 帧 (type 5) 是关键帧，对于 H.265, IDR_W_RADL (type 19) 和 IDR_N_LP (type 20) 是关键帧
    // 在你的代码中，type 19 被认为是 IDR 帧
    return (type == 19 || type == 32 || type == 33 || type == 34); // 假设 VPS, SPS, PPS 也被视为关键帧，用于某些 header 的添加逻辑
}


// --- 主程序 ---
int main() {
    Log("--- 示例程序开始 ---");

    // 模拟更真实的 NALU 数组，包含 VPS, SPS, PPS, IDR 关键帧, 普通帧
    std::vector<Nalu*> nalu_vector_sim;

    // 模拟 VPS, SPS, PPS (通常在 IDR 帧前)
    struct Nalu* vps_nalu = new Nalu();
    vps_nalu->type = 32; // VPS_NUT
    vps_nalu->length = 50;
    vps_nalu->packet.resize(vps_nalu->length);
    std::iota(vps_nalu->packet.begin(), vps_nalu->packet.end(), 1);
    nalu_vector_sim.push_back(vps_nalu);
    out_nalu((char*)vps_nalu->packet.data(), vps_nalu->length, 32);

    struct Nalu* sps_nalu = new Nalu();
    sps_nalu->type = 33; // SPS_NUT
    sps_nalu->length = 100;
    sps_nalu->packet.resize(sps_nalu->length);
    std::iota(sps_nalu->packet.begin(), sps_nalu->packet.end(), 51);
    nalu_vector_sim.push_back(sps_nalu);
    out_nalu((char*)sps_nalu->packet.data(), sps_nalu->length, 33);

    struct Nalu* pps_nalu = new Nalu();
    pps_nalu->type = 34; // PPS_NUT
    pps_nalu->length = 30;
    pps_nalu->packet.resize(pps_nalu->length);
    std::iota(pps_nalu->packet.begin(), pps_nalu->packet.end(), 151);
    nalu_vector_sim.push_back(pps_nalu);
    out_nalu((char*)pps_nalu->packet.data(), pps_nalu->length, 34);

    // 模拟 IDR 关键帧
    struct Nalu* idr_nalu = new Nalu();
    idr_nalu->type = 19; // IDR_W_RADL (示例 IDR 类型)
    idr_nalu->length = 2048;
    idr_nalu->packet.resize(idr_nalu->length);
    std::iota(idr_nalu->packet.begin(), idr_nalu->packet.end(), 201);
    nalu_vector_sim.push_back(idr_nalu);
    out_nalu((char*)idr_nalu->packet.data(), idr_nalu->length, 19);

    // 模拟 普通帧 (非关键帧)
    struct Nalu* non_idr_nalu = new Nalu();
    non_idr_nalu->type = 1; // TRAIL_R_NUT (示例 普通帧类型)
    non_idr_nalu->length = 1500;
    non_idr_nalu->packet.resize(non_idr_nalu->length);
    std::iota(non_idr_nalu->packet.begin(), non_idr_nalu->packet.end(), 2200);
    nalu_vector_sim.push_back(non_idr_nalu);
    out_nalu((char*)non_idr_nalu->packet.data(), non_idr_nalu->length, 1);

    // RTP 发送处理略...
    int time_base = 90000;
    int fps = 24;
    int send_packet_interval = 1000 / fps;
    int interval = time_base / fps;
    long pts = 0;
    int single_packet_max_length = 1400;
    int ssrc_val = 0x12345678; // 示例 SSRC
    std::string rtp_protocol_type = "UDP/RTP/AVP"; // 或 "TCP/RTP/AVP"


    Log("[process_request] 开始模拟推流, NALU 数量: " + std::to_string(nalu_vector_sim.size()));
    for (auto* nalu : nalu_vector_sim) {

        const NaluType type = nalu->type;
        const int length = nalu->length;
        const uint8_t* packet = nalu->packet.data();
        const bool is_key = is_keyframe(type);

        Log("处理 NALU, Type: " + std::to_string(type) + ", Keyframe: " + (is_key ? "Yes" : "No") + ", Size: " + std::to_string(length) + " bytes.");

        // 在遇到 VPS、SPS、PPS 时，先进行封装
        char frame_buffer[1024 * 128]; // 帧数据缓冲区
        int frame_index = 0;
        char ps_header_buf[PS_HDR_LEN];
        char sys_header_buf[SYS_HDR_LEN];
        char psm_header_buf[PSM_HDR_LEN];
        char pes_header_buf[PES_HDR_LEN];
        char rtp_packet_buf[RTP_HDR_LEN + 1400];

        // 声明 rtp_header_buf
        char rtp_header_buf[RTP_HDR_LEN];

        // 封装 VPS, SPS, PPS
        if (type == 32 || type == 33 || type == 34) { // VPS, SPS, PPS
            memcpy(frame_buffer + frame_index, packet, length);
            frame_index += length;
        }

        // --- PS 封装 ---
        gb28181_make_ps_header(ps_header_buf, pts);
        memcpy(frame_buffer + frame_index, ps_header_buf, PS_HDR_LEN);
        frame_index += PS_HDR_LEN;

        if (is_key) {
            gb28181_make_sys_header(sys_header_buf, 0x3f); // 示例 rate_bound
            memcpy(frame_buffer + frame_index, sys_header_buf, SYS_HDR_LEN);
            frame_index += SYS_HDR_LEN;

            gb28181_make_psm_header(psm_header_buf);
            memcpy(frame_buffer + frame_index, psm_header_buf, PSM_HDR_LEN);
            frame_index += PSM_HDR_LEN;
        }

        // --- PES 封装 ---
        gb28181_make_pes_header(pes_header_buf, 0xe0, length, pts, pts);
        memcpy(frame_buffer + frame_index, pes_header_buf, PES_HDR_LEN);
        frame_index += PES_HDR_LEN;

        memcpy(frame_buffer + frame_index, packet, length);
        frame_index += length;


        // --- RTP 分包发送 ---
        int rtp_packet_count = (frame_index + single_packet_max_length - 1) / single_packet_max_length;
        Log("RTP 分包数: " + std::to_string(rtp_packet_count));

        for (int i = 0; i < rtp_packet_count; ++i) {
            bool is_last_packet = (i == rtp_packet_count - 1);
            gb28181_make_rtp_header(rtp_header_buf, rtp_seq_counter, pts, ssrc_val, is_last_packet);

            int offset = i * single_packet_max_length;
            int data_size = std::min(single_packet_max_length, frame_index - offset);
            int rtp_start_index = 0;

            if (rtp_protocol_type == "TCP/RTP/AVP") {
                uint16_t packet_length = RTP_HDR_LEN + data_size;
                rtp_packet_buf[0] = (packet_length >> 8) & 0xFF;
                rtp_packet_buf[1] = packet_length & 0xFF;
                rtp_start_index = 2;
            }

            memcpy(rtp_packet_buf + rtp_start_index, rtp_header_buf, RTP_HDR_LEN);
            memcpy(rtp_packet_buf + rtp_start_index + RTP_HDR_LEN, frame_buffer + offset, data_size);

            send_network_packet(rtp_packet_buf, rtp_start_index + RTP_HDR_LEN + data_size);
            rtp_seq_counter++;
        }

        pts += interval;
        std::this_thread::sleep_for(std::chrono::milliseconds(send_packet_interval));
        delete nalu; // 模拟 Device::push_rtp_stream 中的 nalu delete
    }

    Log("--- 示例程序结束 ---");
    return 0;
}