main.cpp

#include <iostream>
#include <string>
#include <chrono>
#include <thread>
#include <cstring>
#include <ctime>
#include <cstdint>
#include <cctype>
#ifdef _WIN32
#include <windows.h>
#else
#include <unistd.h>
#include <sys/select.h>
#endif

#include "webcam_interface.h"
#include "jpeg_helper.h"
#include "pipeline.h"
#include "video_file_manager.h"
#include "h264_encoder_factory.h"
#include "jpeg_decoder_factory.h"

std::string generateTimestamp() {
    auto now = std::chrono::system_clock::now();
    auto time_t_now = std::chrono::system_clock::to_time_t(now);
    auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        now.time_since_epoch()) % 1000;

    std::tm* tm_info = localtime(&time_t_now);
    char buffer[32];
    strftime(buffer, sizeof(buffer), "%Y%m%d_%H%M%S", tm_info);

    return std::string(buffer) + "_" + std::to_string(ms.count());
}

bool saveAsJpeg(const WebcamFrame& frame, const std::string& filename, int quality = 95) {
    if (frame.empty() || frame.data == nullptr) {
        std::cerr << "Error: Empty frame" << std::endl;
        return false;
    }

    if (frame.isJpeg) {
        FILE* f = fopen(filename.c_str(), "wb");
        if (!f) {
            std::cerr << "Error: Cannot open file " << filename << std::endl;
            return false;
        }
        size_t written = fwrite(frame.data, 1, frame.size, f);
        fclose(f);
        
        if (written == (size_t)frame.size) {
            std::cout << "Saved (raw JPEG): " << filename << std::endl;
            return true;
        } else {
            std::cerr << "Error: Failed to write JPEG data" << std::endl;
            return false;
        }
    } else {
        JpegWriter writer(quality);
        bool result = writer.save(frame.data, frame.width, frame.height, filename.c_str());
        if (result) {
            std::cout << "Saved: " << filename << std::endl;
        }
        return result;
    }
}

void sleep_ms(int milliseconds) {
#ifdef _WIN32
    Sleep(milliseconds);
#else
    usleep(milliseconds * 1000);
#endif
}

bool checkKey(char key) {
    char upperKey = std::toupper(static_cast<unsigned char>(key));
#ifdef _WIN32
    return (GetAsyncKeyState(upperKey) & 0x8000) != 0;
#else
    static char bufferedKey = 0;
    static bool buffered = false;

    fd_set rfds;
    struct timeval tv;
    FD_ZERO(&rfds);
    FD_SET(STDIN_FILENO, &rfds);
    tv.tv_sec = 0;
    tv.tv_usec = 10000;

    if (select(STDIN_FILENO + 1, &rfds, NULL, NULL, &tv) > 0) {
        char readKey;
        if ((read(STDIN_FILENO, &readKey, 1) > 0) && (readKey != '\n')) {
            bufferedKey = readKey;
            buffered = true;
        }
    }

    if (buffered && std::toupper(static_cast<unsigned char>(bufferedKey)) == upperKey) {
        buffered = false;
        return true;
    }
    return false;
#endif
}

void printUsage(const char* programName) {
    std::cout << "Usage: " << programName << " [options]" << std::endl;
    std::cout << "Options:" << std::endl;
#ifdef _WIN32
    std::cout << "  --vfw              Use VFW capture backend (default: DirectShow)" << std::endl;
#endif
    std::cout << "  --test-pattern     Use virtual test-pattern generator (no camera needed)" << std::endl;
    std::cout << "  --test-audio-wave  Use virtual audio wave generator (no audio needed)" << std::endl;
    std::cout << "  --device <index>   Camera device index (default: 0)" << std::endl;
    std::cout << "  --jpeg             Start in JPEG capture mode (default: RTSP)" << std::endl;
    std::cout << "  --rtsp             Start RTSP streaming server (default)" << std::endl;
    std::cout << "  --port <port>      RTSP server port (default: 8554)" << std::endl;
    std::cout << "  --width <width>    Video width (default: 1280)" << std::endl;
    std::cout << "  --height <height>  Video height (default: 720)" << std::endl;
    std::cout << "  --fps <fps>        Video fps (default: 30)" << std::endl;
    std::cout << "  --bitrate <bps>    Video bitrate (default: 2000000)" << std::endl;
    std::cout << "  --encoder <enc>    Encoder type: x264, rockchip (default: x264)" << std::endl;
    std::cout << "  --decoder <dec>    JPEG decoder type: libyuv, rockchip (default: libyuv)" << std::endl;
    std::cout << "  --no-osd           Disable OSD overlay (default: enabled)" << std::endl;
    std::cout << "  --detect           Enable face detection (default: disabled)" << std::endl;
    std::cout << "  --record [file]    Enable recording (optional filename)" << std::endl;
    std::cout << "  --format <fmt>     Recording format: mp4, ts or aac (default: mp4)" << std::endl;
    std::cout << "  --segment-duration <sec>  Max duration per file (default: 300)" << std::endl;
    std::cout << "  --segment-size <MB>       Max size per file in MB (default: 1024)" << std::endl;
    std::cout << "  --max-files <num>         Max files in loop mode (default: 100)" << std::endl;
    std::cout << "  --loop                    Enable loop overwrite mode" << std::endl;
    std::cout << "  --min-disk-space <MB>     Min free disk space in MB (default: 500)" << std::endl;
    std::cout << "  --audio                   Enable audio capture (default: disabled)" << std::endl;
    std::cout << "  --audio-device <idx>      Audio device index: 0,1,2,... (default: auto)" << std::endl;
    std::cout << "  --audio-sample-rate <hz>  Audio sample rate (default: 16000)" << std::endl;
    std::cout << "  --audio-channels <n>      Audio channels: 1=mono, 2=stereo (default: 1)" << std::endl;
    std::cout << "  --audio-bitrate <bps>     Audio bitrate (default: 128000)" << std::endl;
    std::cout << "  --help                    Show this help" << std::endl;
    std::cout << std::endl;
    std::cout << "Default mode: RTSP streaming mode" << std::endl;
    std::cout << "Press 's' to save snapshot, 'q' to quit" << std::endl;
}

bool runJpegMode(const WebcamConfig& config) {
    std::cout << "=== Webcam Capture Program (JPEG Mode) ===" << std::endl;

    if (config.backend == CaptureBackend::TEST_PATTERN) {
        std::cout << "Video capturer: Test Pattern" << std::endl;
    }
    else {
#ifdef _WIN32
        if (config.backend == CaptureBackend::DIRECTSHOW) {
            std::cout << "Platform: Windows (DirectShow)" << std::endl;
        }
        else {
            std::cout << "Platform: Windows (VFW)" << std::endl;
        }
#else
        std::cout << "Platform: Linux (V4L2)" << std::endl;
#endif
    }

    std::cout << "Available devices: " << WebcamCapture::getDeviceCount(config.backend) << std::endl;

    std::string deviceName, devicePath;
    if (WebcamCapture::getDeviceInfo(config.backend, config.deviceIndex, deviceName, devicePath)) {
        std::cout << "Using device: " << deviceName << std::endl;
    }

    std::cout << "Initializing camera..." << std::endl;

    WebcamCapture webcam(config);

    if (!webcam.isOpened()) {
        std::cerr << "Error: Cannot open camera!" << std::endl;
        return false;
    }

    webcam.printInfo();

    std::cout << "\nStarting capture..." << std::endl;
    std::cout << "Press 's' to save snapshot, 'q' to quit" << std::endl;

    int frameCount = 0;
    const int saveInterval = 30;

    while (true) {
        WebcamFrame frame = webcam.read();

        if (frame.empty()) {
            std::cerr << "Error: Failed to capture frame!" << std::endl;
            continue;
        }

        frameCount++;

        if (checkKey('S')) {
            std::string filename = "snapshot_" + generateTimestamp() + ".jpg";
            saveAsJpeg(frame, filename, 95);
        } else if (checkKey('Q')) {
            std::cout << "\nQuitting..." << std::endl;
            break;
        }

        if (frameCount % saveInterval == 0) {
            std::string filename = "auto_capture_" + generateTimestamp() + ".jpg";
            saveAsJpeg(frame, filename, 90);
        }

        frame.release();
        sleep_ms(33);
    }

    return true;
}

bool runRtspMode(const VideoPipeline::Config& config) {
    std::cout << "=== Webcam Capture Program (RTSP Mode) ===" << std::endl;

    if (config.backend == CaptureBackend::TEST_PATTERN) {
        std::cout << "Video capturer: Test Pattern" << std::endl;
    } else {
#ifdef _WIN32
        if (config.backend == CaptureBackend::DIRECTSHOW) {
          std::cout << "Platform: Windows (DirectShow)" << std::endl;
        } else {
          std::cout << "Platform: Windows (VFW)" << std::endl;
        }
#else
        std::cout << "Platform: Linux (V4L2)" << std::endl;
#endif
    }
    if (config.audio.backend == AudioCaptureBackend::TEST_WAVE) {
        std::cout << "Audio capturer: Test wave (single tone signal)" << std::endl;
    }

    std::cout << "Resolution: " << config.width << "x" << config.height << std::endl;
    std::cout << "FPS: " << config.fps << std::endl;
    std::cout << "Bitrate: " << config.bitrate << " bps" << std::endl;
    std::cout << "Encoder: " << (config.encoderType == EncoderType::ROCKCHIP_MPP ? "Rockchip MPP" : "x264") << std::endl;
    std::cout << "JPEG Decoder: " << (config.jpegDecoderType == JpegDecoderType::ROCKCHIP_MPP ? "Rockchip MPP" : "libyuv") << std::endl;
    std::cout << "RTSP Port: " << config.rtspPort << std::endl;
    if (config.audio.enable) {
        std::cout << "Audio: " << config.audio.sampleRate << "Hz, "
                  << config.audio.channels << " channel(s), "
                  << config.audio.bitrate << " bps (FDK-AAC)" << std::endl;
    }
    std::cout << std::endl;

    VideoPipeline pipeline;

    if (!pipeline.start(config)) {
        std::cerr << "Failed to start pipeline" << std::endl;
        return false;
    }

    std::thread rtspThread([&pipeline]() {
        pipeline.getRtspServer()->run();
    });

    std::cout << "RTSP server is running..." << std::endl;
    std::cout << "Press 'q' to quit" << std::endl;

    auto lastStatsTime = std::chrono::steady_clock::now();

    while (pipeline.isRunning()) {
        if (checkKey('Q')) {
            std::cout << "\nQuitting..." << std::endl;
            break;
        }
        if (pipeline.wasDeviceLost()) {
            std::cerr << "Pipeline exited due to device loss" << std::endl;
            break;
        }

        auto now = std::chrono::steady_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - lastStatsTime).count();
        if (elapsed >= 3) {
            auto stats = pipeline.getStats();
            std::cout << "\rFPS - Capt:" << stats.captureFps
                      << " Conv:" << stats.convertFps
                      << " Det:" << stats.detectFps
                      << " OSD:" << stats.osdFps
                      << " Enc:" << stats.encodeFps
                      << " | Queue:" << stats.captureQueueSize
                      << "/" << stats.convertQueueSize
                      << "/" << stats.detectQueueSize
                      << "/" << stats.osdQueueSize
                      << "/" << stats.h264QueueSize;
            if (config.audio.enable) {
                std::cout << " | Aud:" << stats.audioCaptureFps
                          << "/" << stats.audioEncodeFps
                          << " (q:" << stats.audioQueueSize << ")";
            }
            std::cout << " | Drop(enc/stream):" << stats.encodeDrop
                      << "/" << stats.streamDrop
                      << " Deli:" << stats.streamDelivered // Delivered
                      << "    " << std::flush;
            lastStatsTime = now;
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }

    pipeline.stop();

    if (rtspThread.joinable()) {
        rtspThread.join();
    }

    std::cout << "\nPipeline stopped." << std::endl;
    return true;
}

int main(int argc, char* argv[]) {
    bool rtspMode = true;
    int deviceIndex = 0;
    int width = 1280;
    int height = 720;
    int fps = 30;
    int bitrate = 2000000;
    EncoderType encoderType = EncoderType::X264;
    JpegDecoderType jpegDecoderType = JpegDecoderType::LIBYUV;
    int port = 8554;
    bool enableOsd = true;
    bool enableDetect = false;
    bool enableRecord = false;
    WriterFormat recordFormat = WriterFormat::MP4;
    std::string recordFilename;
    bool enableFileManagement = false;
    VideoFileManager::SegmentConfig segmentConfig;
    VideoFileManager::DiskConfig diskConfig;
    VideoFileManager::LoopConfig loopConfig;
    AudioPipelineConfig audioConfig;  // default: enable=false
#ifdef _WIN32
    CaptureBackend backend = CaptureBackend::DIRECTSHOW;
#else
    CaptureBackend backend = CaptureBackend::AUTO;
#endif

    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--help" || arg == "-h") {
            printUsage(argv[0]);
            return 0;
        } else if (arg == "--device" && i + 1 < argc) {
            deviceIndex = std::stoi(argv[++i]);
        } else if (arg == "--jpeg") {
            rtspMode = false;
        } else if (arg == "--rtsp") {
            rtspMode = true;
        } else if (arg == "--port" && i + 1 < argc) {
            port = std::stoi(argv[++i]);
        } else if (arg == "--width" && i + 1 < argc) {
            width = std::stoi(argv[++i]);
        } else if (arg == "--height" && i + 1 < argc) {
            height = std::stoi(argv[++i]);
        } else if (arg == "--fps" && i + 1 < argc) {
            fps = std::stoi(argv[++i]);
        } else if (arg == "--bitrate" && i + 1 < argc) {
            bitrate = std::stoi(argv[++i]);
        } else if (arg == "--encoder" && i + 1 < argc) {
            std::string enc = argv[++i];
            if (enc == "rockchip" || enc == "rockchip_mpp") {
#ifdef __linux__
                encoderType = EncoderType::ROCKCHIP_MPP;
#else
                std::cerr << "Warning: rockchip encoder is only available on Linux. Using x264." << std::endl;
#endif
            } else {
                encoderType = EncoderType::X264;
            }
        } else if (arg == "--decoder" && i + 1 < argc) {
            std::string dec = argv[++i];
            if (dec == "rockchip" || dec == "rockchip_mpp") {
#ifdef __linux__
                jpegDecoderType = JpegDecoderType::ROCKCHIP_MPP;
#else
                std::cerr << "Warning: rockchip decoder is only available on Linux. Using libyuv." << std::endl;
#endif
            } else {
                jpegDecoderType = JpegDecoderType::LIBYUV;
            }
        } else if (arg == "--no-osd") {
            enableOsd = false;
        } else if (arg == "--detect") {
            enableDetect = true;
        }
        else if (arg == "--record") {
            enableRecord = true;
            if (i + 1 < argc && argv[i + 1][0] != '-') {
                recordFilename = argv[++i];
            }
        } else if (arg == "--format" && i + 1 < argc) {
            std::string fmt = argv[++i];
            if (fmt == "ts" || fmt == "mpegts") {
                recordFormat = WriterFormat::MPEGTS;
            } else if (fmt == "aac") {
                recordFormat = WriterFormat::AAC;
            } else {
                recordFormat = WriterFormat::MP4;
            }
        } else if (arg == "--segment-duration" && i + 1 < argc) {
            segmentConfig.maxDurationSec = std::stoi(argv[++i]);
            enableFileManagement = true;
        } else if (arg == "--segment-size" && i + 1 < argc) {
            segmentConfig.maxFileSizeMB = std::stoi(argv[++i]);
            enableFileManagement = true;
        } else if (arg == "--max-files" && i + 1 < argc) {
            loopConfig.maxTotalFiles = std::stoi(argv[++i]);
            enableFileManagement = true;
        } else if (arg == "--loop") {
            loopConfig.enabled = true;
            enableFileManagement = true;
        } else if (arg == "--min-disk-space" && i + 1 < argc) {
            diskConfig.minFreeSpaceMB = std::stoi(argv[++i]);
            enableFileManagement = true;
        } else if (arg == "--audio") {
            audioConfig.enable = true;
        } else if (arg == "--audio-device" && i + 1 < argc) {
            audioConfig.deviceIndex = std::stoi(argv[++i]);
        } else if (arg == "--audio-sample-rate" && i + 1 < argc) {
            audioConfig.sampleRate = std::stoi(argv[++i]);
        } else if (arg == "--audio-channels" && i + 1 < argc) {
            audioConfig.channels = std::stoi(argv[++i]);
        } else if (arg == "--audio-bitrate" && i + 1 < argc) {
            audioConfig.bitrate = std::stoi(argv[++i]);
        } else if (arg == "--vfw") {
#ifdef _WIN32
            backend = CaptureBackend::VFW;
#endif
        } else if (arg == "--test-pattern") {
            backend = CaptureBackend::TEST_PATTERN;
        } else if (arg == "--test-audio-wave") {
            audioConfig.backend = AudioCaptureBackend::TEST_WAVE;
        }
    }

    if (recordFilename == "") {
        if (enableFileManagement) {
            recordFilename = "record";
        } else {
            recordFilename = "record_" + generateTimestamp();
        }
    }

    if (rtspMode) {
        VideoPipeline::Config config;
        config.backend = backend;
        config.deviceIndex = deviceIndex;
        config.width = width;
        config.height = height;
        config.pixelFormat = PixelFormat::MJPG;
        config.fps = fps;
        config.bitrate = bitrate;
        config.encoderType = encoderType;
        config.jpegDecoderType = jpegDecoderType;
        config.rtspPort = port;
        config.enableOsd = enableOsd;
        config.enableDetect = enableDetect;
        config.enableRecord = enableRecord;
        config.recordFormat = recordFormat;
        config.recordFilename = recordFilename;
        config.enableFileManagement = enableFileManagement;
        config.segmentConfig = segmentConfig;
        config.diskConfig = diskConfig;
        config.loopConfig = loopConfig;
        config.audio = audioConfig;
        return runRtspMode(config) ? 0 : -1;
    } else {
        WebcamConfig config;
        config.backend = backend;
        config.deviceIndex = deviceIndex;
        config.width = width;
        config.height = height;
        config.format = PixelFormat::MJPG;
        config.fps = fps;
        return runJpegMode(config) ? 0 : -1;
    }
}