#include "config.hpp"
#include "bytetrack.hpp"
#include "batch_store.hpp"
#include "csv_logger.hpp"
#include "rknn_yolo.hpp"
#include "drawing.hpp"
#include "nms.hpp"

#include <opencv2/opencv.hpp>
#include <Eigen/Dense>
#include <csignal>
#include <iostream>
#include <chrono>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <deque>
#include <ctime>
#include <filesystem>

static std::atomic<bool> shutdown_requested{false};

static void signal_handler(int) {
    shutdown_requested = true;
    std::cout << "\nShutdown requested — finishing current frame..." << std::endl;
}

static std::string now_str() {
    auto now = std::chrono::system_clock::now();
    auto t = std::chrono::system_clock::to_time_t(now);
    std::ostringstream ss;
    ss << std::put_time(std::localtime(&t), "%Y-%m-%d %H:%M:%S");
    return ss.str();
}

static cv::VideoCapture open_capture(const std::string& source) {
    bool is_live = source.substr(0, 7) == "rtsp://" || source.substr(0, 7) == "http://";
    if (is_live) {
        setenv("OPENCV_FFMPEG_CAPTURE_OPTIONS", g_config.rtsp_ffmpeg_options.c_str(), 1);
    }
    cv::VideoCapture cap(source, cv::CAP_FFMPEG);
    cap.set(cv::CAP_PROP_BUFFERSIZE, 1);
    return cap;
}

static void warmup_stream(cv::VideoCapture& cap, int n) {
    std::cout << "Warming up stream..." << std::endl;
    cv::Mat frame;
    for (int i = 0; i < n; ++i) cap.read(frame);
    std::cout << "Stream ready!" << std::endl;
}

struct TrackedInfo {
    float cx;
    std::chrono::steady_clock::time_point ts;
};

static void prune_stale_tracks(std::unordered_map<int, TrackedInfo>& tracked, double prune_sec) {
    auto now = std::chrono::steady_clock::now();
    std::vector<int> stale;
    for (auto& [tid, info] : tracked) {
        auto age = std::chrono::duration_cast<std::chrono::seconds>(now - info.ts).count();
        if (age > prune_sec) stale.push_back(tid);
    }
    for (int tid : stale) tracked.erase(tid);
}

int main(int argc, char* argv[]) {
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    const char* env_path = argc > 1 ? argv[1] : nullptr;
    g_config.load_from_env(env_path);

    std::cout << "Starting ByteTrack Counter (C++)" << std::endl;
    std::cout << "Camera: " << g_config.camera_name << std::endl;
    std::cout << "Source: " << g_config.source << std::endl;
    std::cout << "Model: " << g_config.model_path << std::endl;
    std::cout << "DB: " << g_config.db_path << std::endl;

    int class_ayam = 0;
    int class_talenan = 1;

    auto logger = [](const std::string& msg) {
        std::cout << "[" << now_str() << "] " << msg << std::endl;
    };

    BatchStore store(g_config.db_path, g_config.state_file,
                     g_config.camera_name, g_config.object_label,
                     g_config.daily_cutoff_time,
                     g_config.batch_timeout_seconds,
                     g_config.ignore_batch_label_timeout,
                     g_config.min_object_per_batch,
                     g_config.min_duration_per_batch,
                     logger);
    store.start_cutoff_watcher();

    std::unique_ptr<CsvLogger> cross_logger;
    if (g_config.export_csv) {
        cross_logger = std::make_unique<CsvLogger>(g_config.cross_csv,
            std::vector<std::string>{"batch", "frame", "timestamp", "chicken_id"});
    }

    RKNNYOLO model(g_config.model_path, g_config.core_mask,
                   g_config.imgsz, g_config.conf, 0.45f,
                   g_config.num_classes, 0, g_config.score_sigmoid);

    ByteTracker ayam_tracker(g_config.track_high_thresh, g_config.track_low_thresh,
                              g_config.track_match_thresh, g_config.track_buffer,
                              g_config.track_min_hits);
    ByteTracker talenan_tracker(g_config.track_high_thresh, g_config.track_low_thresh,
                                 g_config.track_match_thresh, g_config.track_buffer,
                                 g_config.track_min_hits);

    std::unordered_set<int> ayam_line_crossed;
    std::unordered_set<int> talenan_line_crossed;

    std::unordered_map<int, int> ayam_cross_flash;
    std::unordered_map<int, int> talenan_cross_flash;
    int line_pulse = 0, count_pulse = 0, batch_pulse = 0;
    std::vector<Popup> popups;

    auto session_start = std::chrono::steady_clock::now();
    int frame_idx = 0;

    cv::VideoCapture cap = open_capture(g_config.source);
    bool is_live = g_config.source.substr(0, 7) == "rtsp://" ||
                   g_config.source.substr(0, 7) == "http://";

    if (!cap.isOpened()) {
        std::cerr << "Cannot open source: " << g_config.source << std::endl;
        store.shutdown();
        return 1;
    }
    if (is_live && g_config.warmup_frames > 0)
        warmup_stream(cap, g_config.warmup_frames);

    int w = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
    int h = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
    double fps_prop = cap.get(cv::CAP_PROP_FPS);
    if (fps_prop <= 1) fps_prop = g_config.output_fps;

    int line_x = resolve_line_x(w, g_config.line_x, g_config.line_x_frac);

    std::cout << "RKNN+ByteTrack counter | " << w << "x" << h << " @ "
              << fps_prop << "fps | line x=" << line_x
              << " | cross=" << g_config.cross_direction << std::endl;

    int reconnect_count = 0;

    std::unordered_map<int, TrackedInfo> ayam_tracked;
    std::unordered_map<int, TrackedInfo> talenan_tracked;

    while (!shutdown_requested) {
        cv::Mat frame;
        cap >> frame;
        if (frame.empty()) {
            if (!is_live) break;
            reconnect_count++;
            std::cout << "Stream dropped (attempt " << reconnect_count
                      << "), reconnecting..." << std::endl;
            cap.release();
            std::this_thread::sleep_for(std::chrono::seconds(g_config.reconnect_delay_sec));
            cap = open_capture(g_config.source);
            if (!cap.isOpened()) break;
            w = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
            h = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
            line_x = resolve_line_x(w, g_config.line_x, g_config.line_x_frac);
            continue;
        }

        auto elapsed_sec = std::chrono::duration_cast<std::chrono::seconds>(
            std::chrono::steady_clock::now() - session_start).count();
        auto mono = std::chrono::steady_clock::now();

        bool ayam_crossed_frame = false, batch_closed_frame = false, batch_started_frame = false;

        std::vector<Detection> detections = model(frame);

        if (!detections.empty()) {
            std::vector<Eigen::Vector4f> ayam_boxes, talenan_boxes;
            std::vector<float> ayam_scores, talenan_scores;
            std::vector<float> ayam_cx_list, talenan_cx_list;
            std::vector<std::vector<std::vector<cv::Point2f>>> ayam_kpts_list;

            for (const auto& det : detections) {
                const auto& bbox = det.bbox;
                float cx = (bbox(0) + bbox(2)) / 2.0f;

                if (det.cls == class_talenan) {
                    talenan_boxes.push_back(bbox);
                    talenan_scores.push_back(det.score);
                    talenan_cx_list.push_back(cx);
                } else if (det.cls == class_ayam) {
                    ayam_boxes.push_back(bbox);
                    ayam_scores.push_back(det.score);
                    ayam_cx_list.push_back(cx);
                    ayam_kpts_list.push_back(det.keypoints);
                }
            }

            int na = static_cast<int>(ayam_boxes.size());
            int nt = static_cast<int>(talenan_boxes.size());

            Eigen::MatrixXf ayam_mat(na, 4), talenan_mat(nt, 4);
            Eigen::VectorXf ayam_scores_vec(na), talenan_scores_vec(nt);

            for (int i = 0; i < na; ++i) {
                ayam_mat.row(i) = ayam_boxes[i];
                ayam_scores_vec(i) = ayam_scores[i];
            }
            for (int i = 0; i < nt; ++i) {
                talenan_mat.row(i) = talenan_boxes[i];
                talenan_scores_vec(i) = talenan_scores[i];
            }

            // Apply NMS per-class to suppress duplicate boxes
            if (na > 0) {
                auto keep = nms(ayam_mat, ayam_scores_vec);
                int nk = static_cast<int>(keep.size());
                Eigen::MatrixXf fm(nk, 4);
                Eigen::VectorXf fs(nk);
                std::vector<Eigen::Vector4f> fb;
                std::vector<float> fcx;
                std::vector<std::vector<std::vector<cv::Point2f>>> fkpts;
                for (int ki = 0; ki < nk; ++ki) {
                    int idx = keep[ki];
                    fm.row(ki) = ayam_mat.row(idx);
                    fs(ki) = ayam_scores_vec(idx);
                    fb.push_back(ayam_boxes[idx]);
                    fcx.push_back(ayam_cx_list[idx]);
                    if (idx < (int)ayam_kpts_list.size())
                        fkpts.push_back(ayam_kpts_list[idx]);
                }
                na = nk;
                ayam_mat = std::move(fm);
                ayam_scores_vec = std::move(fs);
                ayam_boxes = std::move(fb);
                ayam_cx_list = std::move(fcx);
                ayam_kpts_list = std::move(fkpts);
            }
            if (nt > 0) {
                auto keep = nms(talenan_mat, talenan_scores_vec);
                int nk = static_cast<int>(keep.size());
                Eigen::MatrixXf fm(nk, 4);
                Eigen::VectorXf fs(nk);
                std::vector<Eigen::Vector4f> fb;
                std::vector<float> fcx;
                for (int ki = 0; ki < nk; ++ki) {
                    int idx = keep[ki];
                    fm.row(ki) = talenan_mat.row(idx);
                    fs(ki) = talenan_scores_vec(idx);
                    fb.push_back(talenan_boxes[idx]);
                    fcx.push_back(talenan_cx_list[idx]);
                }
                nt = nk;
                talenan_mat = std::move(fm);
                talenan_scores_vec = std::move(fs);
                talenan_boxes = std::move(fb);
                talenan_cx_list = std::move(fcx);
            }

            TrackResult ayam_tr = ayam_tracker.update(ayam_mat, ayam_scores_vec);
            TrackResult talenan_tr = talenan_tracker.update(talenan_mat, talenan_scores_vec);

            // Process talenan crossings
            for (int di = 0; di < nt; ++di) {
                auto it = talenan_tr.det_to_track.find(di);
                if (it == talenan_tr.det_to_track.end()) continue;
                int tid = it->second;
                float cx = talenan_cx_list[di];

                auto tit = talenan_tracked.find(tid);
                if (tit != talenan_tracked.end()) {
                    float prev_cx = tit->second.cx;
                    if (crossed_line(prev_cx, cx, line_x, g_config.cross_direction) &&
                        talenan_line_crossed.find(tid) == talenan_line_crossed.end()) {
                        talenan_line_crossed.insert(tid);
                        if (store.record_talenan_crossing(tid)) {
                            batch_closed_frame = true;
                            talenan_cross_flash[tid] = g_config.cross_flash_frames;
                            popups.push_back({
                                static_cast<int>(cx) - 20,
                                static_cast<int>((talenan_boxes[di](1) + talenan_boxes[di](3)) / 2),
                                frame_idx, "BATCH CLOSED"
                            });
                        }
                    }
                }
                talenan_tracked[tid] = {cx, mono};
            }

            // Process ayam crossings
            for (int di = 0; di < na; ++di) {
                auto it = ayam_tr.det_to_track.find(di);
                if (it == ayam_tr.det_to_track.end()) continue;
                int tid = it->second;
                float cx = ayam_cx_list[di];

                auto tit = ayam_tracked.find(tid);
                if (tit != ayam_tracked.end()) {
                    float prev_cx = tit->second.cx;
                    if (crossed_line(prev_cx, cx, line_x, g_config.cross_direction) &&
                        ayam_line_crossed.find(tid) == ayam_line_crossed.end()) {
                        ayam_line_crossed.insert(tid);
                        auto [count, started_new] = store.record_ayam_crossing(tid);
                        if (cross_logger) {
                            std::ostringstream frame_str, tid_str;
                            frame_str << frame_idx;
                            tid_str << tid;
                            cross_logger->write_row({
                                std::to_string(store.current_batch_number()),
                                frame_str.str(),
                                now_str(),
                                tid_str.str()
                            });
                        }
                        ayam_crossed_frame = true;
                        if (started_new) batch_started_frame = true;
                        ayam_cross_flash[tid] = g_config.cross_flash_frames;
                        popups.push_back({
                            static_cast<int>(cx) - 12,
                            static_cast<int>((ayam_boxes[di](1) + ayam_boxes[di](3)) / 2),
                            frame_idx, "+1"
                        });
                    }
                }
                ayam_tracked[tid] = {cx, mono};
            }

            // Track lost tracks for continuity
            for (auto& [tid, cx] : ayam_tr.lost_map) {
                if (ayam_tracked.find(tid) == ayam_tracked.end())
                    ayam_tracked[tid] = {cx, mono};
            }

            // Draw talenan
            for (int di = 0; di < nt; ++di) {
                auto it = talenan_tr.det_to_track.find(di);
                if (it == talenan_tr.det_to_track.end()) continue;
                int tid = it->second;
                const auto& bbox = talenan_boxes[di];
                int x1 = static_cast<int>(bbox(0)), y1 = static_cast<int>(bbox(1));
                int x2 = static_cast<int>(bbox(2)), y2 = static_cast<int>(bbox(3));
                int flash = talenan_cross_flash.count(tid) ? talenan_cross_flash[tid] : 0;
                cv::Scalar color = flash > 0 ? C_GREEN : C_TALENAN_BOX;
                cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, flash > 0 ? 3 : 2);
                draw_pill(frame, "TALENAN " + std::to_string(tid), x1, y1 - 4, color);
            }

            // Draw ayam
            for (int di = 0; di < na; ++di) {
                auto it = ayam_tr.det_to_track.find(di);
                if (it == ayam_tr.det_to_track.end()) continue;
                int tid = it->second;
                const auto& bbox = ayam_boxes[di];
                int x1 = static_cast<int>(bbox(0)), y1 = static_cast<int>(bbox(1));
                int x2 = static_cast<int>(bbox(2)), y2 = static_cast<int>(bbox(3));
                int flash = ayam_cross_flash.count(tid) ? ayam_cross_flash[tid] : 0;
                cv::Scalar color = flash > 0 ? C_GREEN : C_AYAM_BOX;
                cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, flash > 0 ? 3 : 2);
                draw_pill(frame, "ID " + std::to_string(tid), x1, y1 - 4, color);
                if (di < (int)ayam_kpts_list.size() && !ayam_kpts_list[di].empty())
                    draw_skeleton_bold(frame, ayam_kpts_list[di]);
            }
        }

        if (ayam_crossed_frame) {
            line_pulse = g_config.line_pulse_frames;
            count_pulse = g_config.count_pulse_frames;
        }
        if (batch_closed_frame) {
            line_pulse = g_config.line_pulse_frames;
        }
        if (batch_started_frame) {
            batch_pulse = g_config.batch_pulse_frames;
        }

        int batch_num = store.current_batch_number();
        int batch_count = store.current_batch_count();
        int display_total = store.display_total();
        double elapsed = static_cast<double>(elapsed_sec);
        double rate = elapsed > 0 ? (display_total / elapsed * 60.0) : 0.0;

        draw_elegant_counting_line(frame, line_x, h, line_pulse);
        draw_hero_count(frame, line_x, h, batch_count, count_pulse);
        draw_hud(frame, w, batch_num, batch_count, display_total,
                 elapsed, rate, g_config.camera_name, now_str());
        draw_batch_banner(frame, w, batch_num, batch_pulse);
        draw_footer(frame, w, h, frame_idx, is_live ? "LIVE-RKNN-BT" : "FILE-RKNN-BT");
        popups = draw_popups(frame, popups, frame_idx);

        // Flash decay
        for (auto& flash_map : {&ayam_cross_flash, &talenan_cross_flash}) {
            for (auto it = flash_map->begin(); it != flash_map->end(); ) {
                it->second--;
                if (it->second <= 0) it = flash_map->erase(it);
                else ++it;
            }
        }
        line_pulse = std::max(0, line_pulse - 1);
        count_pulse = std::max(0, count_pulse - 1);
        batch_pulse = std::max(0, batch_pulse - 1);

        if (g_config.live_stream_enabled && frame_idx % g_config.live_stream_every_n == 0) {
            try {
                std::filesystem::path p(g_config.live_stream_frame_path);
                if (p.has_parent_path())
                    std::filesystem::create_directories(p.parent_path());
                std::vector<uint8_t> jpeg;
                cv::imencode(".jpg", frame, jpeg,
                    {cv::IMWRITE_JPEG_QUALITY, g_config.live_stream_quality});
                std::ofstream f(g_config.live_stream_frame_path, std::ios::binary);
                f.write(reinterpret_cast<const char*>(jpeg.data()), jpeg.size());
            } catch (...) {}
        }

        frame_idx++;
        if (frame_idx % g_config.flush_every_n_frames == 0) {
            std::cout << "[" << now_str() << "] Frame " << frame_idx
                      << " | Batch " << batch_num << ": " << batch_count
                      << " | Total: " << display_total
                      << " | Uptime " << (elapsed / 3600.0) << "h" << std::endl;
        }

        prune_stale_tracks(ayam_tracked, g_config.tracked_prune_sec);
        prune_stale_tracks(talenan_tracked, g_config.tracked_prune_sec);
    }

    cap.release();
    store.shutdown();

    std::cout << "\n=== Batch Summary (SQLite) ===" << std::endl;
    std::cout << "Database: " << g_config.db_path << std::endl;

    return 0;
}