#include "web_server.h"
#include "httplib.h"
#include "json.hpp"
#include <iostream>
#include <fstream>
#include <sstream>
#include <thread>
#include <chrono>
#include <set>
#include <random>

using namespace std;
using json = nlohmann::json;

class WebServer::Impl {
public:
    httplib::Server svr;
    string host;
    int port;
    std::atomic<bool> stopping{false};   // 退出标记: 让MJPEG推流循环及时结束
    std::atomic<bool> mjpeg_enabled{true}; // MJPEG推流开关
    unordered_map<string, function<string()>> get_handlers;
    unordered_map<string, function<string(const map<string,string>&)>> getq_handlers;
    unordered_map<string, function<string(const string&)>> post_handlers;
    unordered_map<int, function<cv::Mat()>> frame_providers;
    unordered_map<int, function<uint64_t()>> seq_providers;

    AuthHooks hooks_;
    bool hooks_set_ = false;

    // ---- 登录会话 ----
    std::set<std::string> sessions;
    std::mutex sess_mtx;

    std::string makeSession() {
        static std::random_device rd;
        static std::mt19937 gen(rd());
        static std::uniform_int_distribution<int> d(0, 15);
        static const char* hx = "0123456789abcdef";
        std::string t;
        for (int i = 0; i < 40; i++) t += hx[d(gen)];
        std::lock_guard<std::mutex> lk(sess_mtx);
        sessions.insert(t);
        return t;
    }
    bool sessionOk(const httplib::Request& req) {
        std::string c = req.get_header_value("Cookie");
        size_t p = c.find("za_session=");
        if (p == std::string::npos) return false;
        p += 11;
        size_t e = c.find(';', p);
        std::string tok = c.substr(p, (e == std::string::npos ? c.size() : e) - p);
        std::lock_guard<std::mutex> lk(sess_mtx);
        return sessions.count(tok) > 0;
    }
    void dropSession(const httplib::Request& req) {
        std::string c = req.get_header_value("Cookie");
        size_t p = c.find("za_session=");
        if (p == std::string::npos) return;
        p += 11;
        size_t e = c.find(';', p);
        std::string tok = c.substr(p, (e == std::string::npos ? c.size() : e) - p);
        std::lock_guard<std::mutex> lk(sess_mtx);
        sessions.erase(tok);
    }

    bool authed(const httplib::Request& req) {
        if (!hooks_set_) return true;   // 未配置凭据时不启用登录
        std::string p = req.path;
        if (p == "/login" || p == "/api/login" || p.rfind("/logo.", 0) == 0 ||
            p == "/favicon.ico")
            return true;
        return sessionOk(req);
    }
    void denyJson(httplib::Response& res) {
        res.status = 401;
        res.set_content(R"({"error":"unauthorized"})", "application/json");
    }

    void setupRoutes() {
        svr.Get("/", [this](const httplib::Request& req, httplib::Response& res) {
            if (hooks_set_ && !authed(req)) {
                res.status = 302;
                res.set_header("Location", "/login");
                return;
            }
            ifstream f("web/index.html");
            if (f.is_open()) {
                string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
                res.set_content(content, "text/html");
            } else {
                res.set_content("<h1>zone_alarm</h1><p>web/index.html not found</p>", "text/html");
            }
        });

        svr.Get("/login", [](const httplib::Request& req, httplib::Response& res) {
            ifstream f("web/login.html");
            if (f.is_open()) {
                string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
                res.set_content(content, "text/html");
            } else {
                res.status = 404;
                res.set_content("login page missing", "text/plain");
            }
        });

        svr.Post("/api/login", [this](const httplib::Request& req, httplib::Response& res) {
            try {
                auto j = json::parse(req.body);
                string u = j.value("username", "");
                string p = j.value("password", "");
                if (hooks_set_ && hooks_.check(u, p)) {
                    std::string tok = makeSession();
                    res.set_header("Set-Cookie",
                                   "za_session=" + tok + "; Path=/; HttpOnly");
                    res.set_content(R"({"ok":true})", "application/json");
                    cout << "[Auth] 登录成功: " << u << endl;
                } else {
                    res.set_content(R"({"ok":false,"error":"用户名或密码错误"})",
                                    "application/json");
                }
            } catch (...) {
                res.set_content(R"({"ok":false})", "application/json");
            }
        });

        svr.Post("/api/logout", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            dropSession(req);
            res.set_header("Set-Cookie",
                           "za_session=deleted; Path=/; Max-Age=0");
            res.set_content(R"({"ok":true})", "application/json");
        });

        svr.Get("/api/auth/info", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            string u = (hooks_set_ && hooks_.currentUser) ? hooks_.currentUser() : "";
            json j; j["username"] = u;
            res.set_content(j.dump(), "application/json");
        });

        svr.Post("/api/auth/changepwd", [this](
                                               const httplib::Request& req,
                                               httplib::Response& res) {
            if (!hooks_set_) { denyJson(res); return; }
            if (!authed(req)) { denyJson(res); return; }
            try {
                auto j = json::parse(req.body);
                string oldp = j.value("old_password", "");
                string u = j.value("username", "");
                string newp = j.value("new_password", "");
                if (!hooks_.update(oldp, u, newp)) {
                    res.set_content(R"({"ok":false,"error":"原密码错误或新密码为空"})",
                                    "application/json");
                    return;
                }
                {
                    std::lock_guard<std::mutex> lk(sess_mtx);
                    sessions.clear();   // 凭据变更后作废所有旧会话
                }
                std::string tok = makeSession();   // 给当前浏览器签发新会话
                res.set_header("Set-Cookie",
                               "za_session=" + tok + "; Path=/; HttpOnly");
                res.set_content(R"({"ok":true})", "application/json");
            } catch (...) {
                res.set_content(R"({"ok":false})", "application/json");
            }
        });

        svr.Get("/logo.ico", [](const httplib::Request& req, httplib::Response& res) {
            ifstream f("logo.ico", ios::binary);
            if (!f.is_open()) { res.status = 404; return; }
            string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
            res.set_content(content, "image/x-icon");
        });

        svr.Get("/logo.png", [](const httplib::Request& req, httplib::Response& res) {
            ifstream f("logo.png", ios::binary);
            if (!f.is_open()) { res.status = 404; return; }
            string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
            res.set_content(content, "image/png");
        });

        svr.Get("/preview.html", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { res.status = 302; res.set_header("Location", "/login"); return; }
            ifstream f("web/preview.html");
            if (f.is_open()) {
                string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
                res.set_content(content, "text/html");
            } else {
                res.status = 404;
                res.set_content("preview.html not found", "text/plain");
            }
        });

        svr.Get("/config.html", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { res.status = 302; res.set_header("Location", "/login"); return; }
            ifstream f("web/config.html");
            if (f.is_open()) {
                string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
                res.set_content(content, "text/html");
            } else {
                res.status = 404;
                res.set_content("config.html not found", "text/plain");
            }
        });

        svr.Get("/stream/:ch", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            int ch = 0;
            try { ch = stoi(req.path_params.at("ch")); } catch (...) {}
            auto fp = frame_providers.find(ch);
            if (fp == frame_providers.end() || !fp->second) {
                res.set_content("no stream", "text/plain");
                return;
            }
            cv::Mat frame = fp->second();
            if (frame.empty()) {
                res.set_content("no frame", "text/plain");
                return;
            }
            vector<uchar> buf;
            vector<int> params = {cv::IMWRITE_JPEG_QUALITY, 70};
            cv::imencode(".jpg", frame, buf, params);
            res.set_content(string(reinterpret_cast<char*>(buf.data()), buf.size()),
                            "image/jpeg");
        });

        svr.Get("/mjpeg/:ch", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            if (!mjpeg_enabled.load()) {
                res.set_content("stream disabled", "text/plain");
                return;
            }
            int ch = 0;
            try { ch = stoi(req.path_params.at("ch")); } catch (...) {}
            auto fp = frame_providers.find(ch);
            if (fp == frame_providers.end() || !fp->second) {
                res.set_content("no stream", "text/plain");
                return;
            }
            auto seqf = seq_providers.find(ch);
            function<uint64_t()> seqfn =
                (seqf != seq_providers.end()) ? seqf->second : nullptr;
            res.set_chunked_content_provider(
                "multipart/x-mixed-replace; boundary=mjpegboundary",
                [this, fp, seqfn](size_t, httplib::DataSink& sink) -> bool {
                    uint64_t last_seq = 0;
                    while (!stopping.load() && mjpeg_enabled.load() && sink.is_writable()) {
                        uint64_t seq = seqfn ? seqfn() : 0;
                        if (seq != 0 && seq == last_seq) {
                            // 帧未更新, 跳过重复JPEG编码
                            std::this_thread::sleep_for(std::chrono::milliseconds(15));
                            continue;
                        }
                        cv::Mat frame = fp->second();
                        if (!frame.empty()) {
                            std::vector<uchar> buf;
                            std::vector<int> params = {cv::IMWRITE_JPEG_QUALITY, 60};
                            cv::imencode(".jpg", frame, buf, params);
                            std::string header = "--mjpegboundary\r\n"
                                "Content-Type: image/jpeg\r\n"
                                "Content-Length: " + std::to_string(buf.size()) + "\r\n\r\n";
                            if (!sink.write(header.data(), header.size())) return false;
                            if (!sink.write(reinterpret_cast<char*>(buf.data()), buf.size())) return false;
                            if (!sink.write("\r\n", 2)) return false;
                            last_seq = seq;
                        }
                        std::this_thread::sleep_for(std::chrono::milliseconds(33));
                    }
                    return true;
                });
        });

        svr.Get("/data/screenshots/:file", [](const httplib::Request& req, httplib::Response& res) {
            string filename = req.path_params.at("file");
            if (filename.find("..") != string::npos) {
                res.status = 403;
                return;
            }
            string path = "data/screenshots/" + filename;
            ifstream f(path, ios::binary);
            if (!f.is_open()) {
                res.status = 404;
                return;
            }
            string content((istreambuf_iterator<char>(f)), istreambuf_iterator<char>());
            res.set_content(content, "image/jpeg");
        });

        svr.Get("/api/status", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/status");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Get("/api/fences", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/fences");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/fences", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/fences");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/config", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/config");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/relay", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/relay");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/relay/status", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/relay/status");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/relay/config", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/relay/config");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/relay/mode", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/relay/mode");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/relay/bindings", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/relay/bindings");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/relay/bindings", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/relay/bindings");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/system", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/system");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/reboot", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/reboot");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/channels/save", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/channels/save");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/channels/delete", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/channels/delete");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/channels/test", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/channels/test");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/buckets", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/buckets");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/buckets", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/buckets");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/buckets/status", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/buckets/status");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/buckets/activate", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/buckets/activate");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/buckets/mode", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/buckets/mode");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/module/config", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/module/config");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/module/config", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/module/config");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/alarms", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = getq_handlers.find("/api/alarms");
            if (it != getq_handlers.end()) {
                std::map<std::string, std::string> ps;
                for (auto& p : req.params) ps[p.first] = p.second;
                res.set_content(it->second(ps), "application/json");
            }
        });

        svr.Get("/api/relay/aliases", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/relay/aliases");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/relay/aliases", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/relay/aliases");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/relay/actions", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/relay/actions");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/relay/actions/clear", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/relay/actions/clear");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/alarms/count", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/alarms/count");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/alarms/delete", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/alarms/delete");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Post("/api/alarms/clear", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/alarms/clear");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/pose/status", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/pose/status");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/pose/config", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/pose/config");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });

        svr.Get("/api/mjpeg/status", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = get_handlers.find("/api/mjpeg/status");
            if (it != get_handlers.end())
                res.set_content(it->second(), "application/json");
        });

        svr.Post("/api/mjpeg/toggle", [this](const httplib::Request& req, httplib::Response& res) {
            if (!authed(req)) { denyJson(res); return; }
            auto it = post_handlers.find("/api/mjpeg/toggle");
            if (it != post_handlers.end())
                res.set_content(it->second(req.body), "application/json");
        });
    }
};

WebServer::WebServer() : impl_(new Impl()) {}

void WebServer::start() {
    impl_->host = host_;
    impl_->port = port_;
    impl_->setupRoutes();

    running_ = true;
    server_thread_ = thread([this]() {
        cout << "[Web] 服务启动: http://" << host_ << ":" << port_ << endl;
        impl_->svr.listen(host_, port_);
        running_ = false;
    });
}

void WebServer::stop() {
    if (impl_) {
        impl_->stopping = true;   // 先通知推流线程退出, 避免join卡死
        impl_->svr.stop();
    }
    if (server_thread_.joinable()) server_thread_.join();
    running_ = false;
}

void WebServer::setMjpegEnabled(bool en) {
    if (impl_) impl_->mjpeg_enabled.store(en);
}

bool WebServer::mjpegEnabled() const {
    return impl_ ? impl_->mjpeg_enabled.load() : true;
}

void WebServer::setGetHandler(const string& path, function<string()> handler) {
    impl_->get_handlers[path] = handler;
}

void WebServer::setPostHandler(const string& path, function<string(const string&)> handler) {
    impl_->post_handlers[path] = handler;
}

void WebServer::setFrameProvider(int channel, function<cv::Mat()> provider) {
    if (channel >= 0)
        impl_->frame_providers[channel] = std::move(provider);
}

void WebServer::setFrameSeqProvider(int channel, function<uint64_t()> provider) {
    if (channel >= 0)
        impl_->seq_providers[channel] = std::move(provider);
}

void WebServer::removeFrameProvider(int channel) {
    impl_->frame_providers.erase(channel);
    impl_->seq_providers.erase(channel);
}

void WebServer::setGetHandlerQ(const std::string& path,
        std::function<std::string(const std::map<std::string, std::string>&)> handler) {
    impl_->getq_handlers[path] = std::move(handler);
}

void WebServer::setAuthHooks(AuthHooks hooks) {
    impl_->hooks_ = std::move(hooks);
    impl_->hooks_set_ = true;
}