#include "rtsp_writer.hpp"
#include "../common/logger.hpp"

namespace pipeline {

RtspWriter::RtspWriter() = default;

RtspWriter::~RtspWriter() {
    release();
}

bool RtspWriter::init(const OutputTargetConfig& config) {
    std::lock_guard<std::mutex> lock(mutex_);
    
    if (is_initialized_) {
        Logger::warning("RtspWriter already initialized");
        return true;
    }

    // 基本配置验证
    if (config.name.empty()) {
        setLastError("Empty name in config");
        return false;
    }

    if (config.path.empty()) {
        setLastError("Empty path in config");
        return false;
    }

    if (config.fps <= 0) {
        setLastError("Invalid fps in config: " + std::to_string(config.fps));
        return false;
    }

    name_ = config.name;
    url_ = config.path;
    fps_ = config.fps;
    codec_ = config.codec;
    bitrate_ = config.bitrate;

    // 对于mock URL，只初始化编码器
    if (url_.substr(0, 7) == "mock://") {
        if (!initFFmpeg()) {
            return false;
        }
        is_initialized_ = true;
        Logger::info("RtspWriter initialized successfully: " + name_);
        return true;
    }

    // 对于实际URL，完整初始化
    if (!initFFmpeg()) {
        return false;
    }

    is_initialized_ = true;
    Logger::info("RtspWriter initialized successfully: " + name_);
    return true;
}

bool RtspWriter::write(const cv::Mat& frame) {
    std::lock_guard<std::mutex> lock(mutex_);
    
    if (!is_initialized_) {
        setLastError("RtspWriter not initialized");
        return false;
    }

    if (frame.empty()) {
        setLastError("Input frame is empty");
        return false;
    }

    try {
        // 对于mock URL，只更新状态
        if (url_.substr(0, 7) == "mock://") {
            is_opened_ = true;
            return true;
        }

        // 延迟打开RTSP输出，直到收到第一帧
        if (!format_ctx_) {
            if (!openRtspOutput()) {
                return false;
            }
        }

        // 转换图像格式
        uint8_t* srcData[4] = {frame.data, nullptr, nullptr, nullptr};
        int srcLinesize[4] = {static_cast<int>(frame.step), 0, 0, 0};
        sws_scale(sws_ctx_, srcData, srcLinesize, 0, frame.rows,
                 frame_->data, frame_->linesize);

        frame_->pts += av_rescale_q(1, codec_ctx_->time_base, video_stream_->time_base);

        // 编码并发送
        int ret = avcodec_send_frame(codec_ctx_, frame_);
        if (ret < 0) {
            setLastError("Error sending frame to encoder");
            return false;
        }

        while (ret >= 0) {
            ret = avcodec_receive_packet(codec_ctx_, packet_);
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                break;
            } else if (ret < 0) {
                setLastError("Error receiving packet from encoder");
                return false;
            }

            packet_->stream_index = video_stream_->index;
            ret = av_interleaved_write_frame(format_ctx_, packet_);
            if (ret < 0) {
                setLastError("Error writing packet");
                return false;
            }
        }

        is_opened_ = true;
        return true;
    } catch (const std::exception& e) {
        setLastError("Error writing frame: " + std::string(e.what()));
        return false;
    }
}

bool RtspWriter::isOpened() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return is_initialized_ && is_opened_;
}

void RtspWriter::release() {
    std::lock_guard<std::mutex> lock(mutex_);
    cleanup();
    is_initialized_ = false;
    is_opened_ = false;
}

std::string RtspWriter::getName() const {
    return name_;
}

bool RtspWriter::getStatus(std::string& error_msg) const {
    std::lock_guard<std::mutex> lock(mutex_);
    error_msg = last_error_;
    return is_initialized_ && is_opened_;
}

bool RtspWriter::initFFmpeg() {
    // 初始化编码器
    const AVCodec* codec;
    if (codec_ == "h264") {
        codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    } else if (codec_ == "h265") {
        codec = avcodec_find_encoder(AV_CODEC_ID_H265);
    } else {
        setLastError("Unsupported codec: " + codec_);
        return false;
    }

    if (!codec) {
        setLastError("Codec not found");
        return false;
    }

    codec_ctx_ = avcodec_alloc_context3(codec);
    if (!codec_ctx_) {
        setLastError("Could not allocate codec context");
        return false;
    }

    // 设置编码器参数
    codec_ctx_->bit_rate = bitrate_;
    codec_ctx_->width = 1920;  // 默认分辨率，将在收到第一帧时更新
    codec_ctx_->height = 1080;
    codec_ctx_->time_base = (AVRational){1, fps_};
    codec_ctx_->framerate = (AVRational){fps_, 1};
    codec_ctx_->gop_size = 12;
    codec_ctx_->pix_fmt = AV_PIX_FMT_YUV420P;
    codec_ctx_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    // 对于mock URL，不需要真正打开编码器
    if (url_.substr(0, 7) == "mock://") {
        return true;
    }

    if (avcodec_open2(codec_ctx_, codec, nullptr) < 0) {
        setLastError("Could not open codec");
        cleanup();
        return false;
    }

    return true;
}

bool RtspWriter::openRtspOutput() {
    // 创建输出上下文
    AVDictionary* options = nullptr;
    
    // 设置协议白名单
    av_dict_set(&options, "protocol_whitelist", "file,unix,tcp,rtsp,rtp,udp", 0);
    
    avformat_alloc_output_context2(&format_ctx_, nullptr, "rtsp", url_.c_str());
    if (!format_ctx_) {
        setLastError("Could not create output context");
        av_dict_free(&options);
        return false;
    }

    // 创建视频流
    video_stream_ = avformat_new_stream(format_ctx_, nullptr);
    if (!video_stream_) {
        setLastError("Could not create video stream");
        av_dict_free(&options);
        cleanup();
        return false;
    }

    video_stream_->time_base = codec_ctx_->time_base;
    avcodec_parameters_from_context(video_stream_->codecpar, codec_ctx_);

    // 分配帧和数据包
    frame_ = av_frame_alloc();
    if (!frame_) {
        setLastError("Could not allocate frame");
        av_dict_free(&options);
        cleanup();
        return false;
    }

    frame_->format = codec_ctx_->pix_fmt;
    frame_->width = codec_ctx_->width;
    frame_->height = codec_ctx_->height;
    frame_->pts = 0;

    if (av_frame_get_buffer(frame_, 0) < 0) {
        setLastError("Could not allocate frame data");
        av_dict_free(&options);
        cleanup();
        return false;
    }

    packet_ = av_packet_alloc();
    if (!packet_) {
        setLastError("Could not allocate packet");
        av_dict_free(&options);
        cleanup();
        return false;
    }

    // 创建图像格式转换上下文
    sws_ctx_ = sws_getContext(
        codec_ctx_->width, codec_ctx_->height, AV_PIX_FMT_BGR24,
        codec_ctx_->width, codec_ctx_->height, codec_ctx_->pix_fmt,
        SWS_BICUBIC, nullptr, nullptr, nullptr);
    if (!sws_ctx_) {
        setLastError("Could not create scale context");
        av_dict_free(&options);
        cleanup();
        return false;
    }

    // 设置RTSP选项
    if (url_.substr(0, 7) != "unix://") {
        av_dict_set(&options, "rtsp_transport", "tcp", 0);
    }
    av_dict_set(&options, "muxdelay", "0.1", 0);

    if (avformat_write_header(format_ctx_, &options) < 0) {
        setLastError("Could not write header");
        av_dict_free(&options);
        cleanup();
        return false;
    }

    av_dict_free(&options);
    return true;
}

void RtspWriter::cleanup() {
    if (format_ctx_) {
        av_write_trailer(format_ctx_);
        avio_closep(&format_ctx_->pb);
        avformat_free_context(format_ctx_);
        format_ctx_ = nullptr;
    }
    
    if (codec_ctx_) {
        avcodec_free_context(&codec_ctx_);
    }
    
    if (frame_) {
        av_frame_free(&frame_);
    }
    
    if (packet_) {
        av_packet_free(&packet_);
    }
    
    if (sws_ctx_) {
        sws_freeContext(sws_ctx_);
        sws_ctx_ = nullptr;
    }
}

void RtspWriter::setLastError(const std::string& error) {
    last_error_ = error;
    Logger::error("RtspWriter error: " + error);
}

} // namespace pipeline