#include <algorithm>
#include <cctype>
#include <cstdio>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

#include <curl/curl.h>
#include <nlohmann/json.hpp>

namespace fs = std::filesystem;

namespace {

const char* kDefaultBaseURL = "https://www.watrixgait.com";
const char* kImageExts[] = {".jpg", ".jpeg", ".png", ".bmp", ".webp"};
constexpr uintmax_t kMaxBatchUploadBytes = 32ULL << 20;

struct HTTPResponse {
    long status = 0;
    std::string body;
};

struct SequenceFeature {
    std::string source;
    std::string sequence_id;
    int frame_count = 0;
    std::vector<double> gait;
    std::vector<double> reid;
    std::vector<double> face;
};

static size_t writeCallback(char* ptr, size_t size, size_t nmemb, void* userdata) {
    auto* out = static_cast<std::string*>(userdata);
    out->append(ptr, size * nmemb);
    return size * nmemb;
}

static std::string envOr(const char* name, const std::string& fallback) {
    const char* value = std::getenv(name);
    if (value && *value) return value;
    return fallback;
}

static std::string apiKey() {
    std::string key = envOr("WATRIX_GAIT_API_KEY", "");
    if (key.empty()) throw std::runtime_error("请设置 WATRIX_GAIT_API_KEY");
    return key;
}

static std::string baseURL() {
    std::string url = envOr("WATRIX_GAIT_BASE_URL", kDefaultBaseURL);
    while (!url.empty() && url.back() == '/') url.pop_back();
    return url;
}

static std::string lowerExt(const fs::path& path) {
    std::string ext = path.extension().string();
    std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char ch) {
        return static_cast<char>(std::tolower(ch));
    });
    return ext;
}

static bool isImage(const fs::path& path) {
    std::string ext = lowerExt(path);
    for (const char* item : kImageExts) {
        if (ext == item) return true;
    }
    return false;
}

static std::vector<fs::path> collectFrames(const fs::path& dir) {
    std::vector<fs::path> frames;
    for (const auto& entry : fs::directory_iterator(dir)) {
        if (entry.is_regular_file() && isImage(entry.path())) {
            frames.push_back(entry.path());
        }
    }
    std::sort(frames.begin(), frames.end());
    return frames;
}

static std::vector<fs::path> collectLeafSequenceDirs(const fs::path& root) {
    std::vector<fs::path> dirs;
    if (!fs::exists(root)) return dirs;
    for (const auto& entry : fs::recursive_directory_iterator(root)) {
        if (!entry.is_directory()) continue;
        bool hasChildDir = false;
        bool hasImage = false;
        for (const auto& child : fs::directory_iterator(entry.path())) {
            if (child.is_directory()) hasChildDir = true;
            if (child.is_regular_file() && isImage(child.path())) hasImage = true;
        }
        if (!hasChildDir && hasImage) dirs.push_back(entry.path());
    }
    std::sort(dirs.begin(), dirs.end());
    return dirs;
}

static HTTPResponse requestJSON(const std::string& method, const std::string& path, const nlohmann::json* payload) {
    CURL* curl = curl_easy_init();
    if (!curl) throw std::runtime_error("curl_easy_init failed");

    std::string body;
    std::string url = baseURL() + path;
    std::string payloadText = payload ? payload->dump() : "";
    std::string auth = "Authorization: Bearer " + apiKey();

    struct curl_slist* headers = nullptr;
    headers = curl_slist_append(headers, "Accept: application/json");
    headers = curl_slist_append(headers, auth.c_str());
    if (payload) headers = curl_slist_append(headers, "Content-Type: application/json");

    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_CUSTOMREQUEST, method.c_str());
    curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeCallback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &body);
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 1800L);
    if (payload) {
        curl_easy_setopt(curl, CURLOPT_POSTFIELDS, payloadText.c_str());
        curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, payloadText.size());
    }

    CURLcode code = curl_easy_perform(curl);
    long status = 0;
    curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &status);
    curl_slist_free_all(headers);
    curl_easy_cleanup(curl);
    if (code != CURLE_OK) throw std::runtime_error(curl_easy_strerror(code));
    if (status >= 400) {
        std::ostringstream oss;
        oss << "HTTP " << status << " " << method << " " << path << "\n" << body;
        throw std::runtime_error(oss.str());
    }
    return {status, body};
}

static void uploadBatch(const std::string& uploadPath, const std::string& uploadToken, const std::vector<fs::path>& frames) {
    // The API accepts one multipart batch with a total size of at most 32 MB.
    uintmax_t totalBytes = 0;
    for (const auto& frame : frames) {
        uintmax_t size = fs::file_size(frame);
        totalBytes += size;
        if (totalBytes > kMaxBatchUploadBytes) throw std::runtime_error("sequence frames exceed 32 MB; reduce frames or compress images");
    }
    CURL* curl = curl_easy_init();
    if (!curl) throw std::runtime_error("curl_easy_init failed");
    std::string url = baseURL() + uploadPath;
    std::string responseBody;
    curl_mime* mime = curl_mime_init(curl);
    curl_mimepart* tokenPart = curl_mime_addpart(mime);
    curl_mime_name(tokenPart, "upload_token");
    curl_mime_data(tokenPart, uploadToken.c_str(), CURL_ZERO_TERMINATED);
    for (const auto& frame : frames) {
        // Repeated frames fields preserve frame order; the API ignores filenames.
        curl_mimepart* framePart = curl_mime_addpart(mime);
        curl_mime_name(framePart, "frames");
        curl_mime_filedata(framePart, frame.string().c_str());
    }
    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_MIMEPOST, mime);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeCallback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &responseBody);
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 1800L);

    CURLcode code = curl_easy_perform(curl);
    long status = 0;
    curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &status);
    curl_mime_free(mime);
    curl_easy_cleanup(curl);
    if (code != CURLE_OK) throw std::runtime_error(curl_easy_strerror(code));
    if (status >= 400) {
        std::ostringstream oss;
        oss << "HTTP " << status << " POST " << uploadPath << "\n" << responseBody;
        throw std::runtime_error(oss.str());
    }
}

static void writeJSON(const fs::path& path, const nlohmann::json& data) {
    fs::create_directories(path.parent_path());
    std::ofstream out(path);
    out << data.dump(2) << "\n";
}

static std::vector<double> numericVector(const nlohmann::json& value) {
    std::vector<double> out;
    if (!value.is_array()) return out;
    for (const auto& item : value) {
        if (item.is_number()) out.push_back(item.get<double>());
    }
    return out;
}

static nlohmann::json dotProduct(const std::vector<double>& left, const std::vector<double>& right) {
    if (left.empty() || right.empty()) {
        return {{"score", nullptr}, {"reason", "empty_feature"}, {"left_dim", left.size()}, {"right_dim", right.size()}};
    }
    size_t used = std::min(left.size(), right.size());
    double score = 0.0;
    for (size_t i = 0; i < used; ++i) score += left[i] * right[i];
    return {
        {"score", score},
        {"left_dim", left.size()},
        {"right_dim", right.size()},
        {"used_dim", used},
        {"dimension_mismatch", left.size() != right.size()},
    };
}

static nlohmann::json identity(const SequenceFeature& item) {
    return {
        {"source", item.source},
        {"sequence_id", item.sequence_id},
        {"frame_count", item.frame_count},
    };
}

static nlohmann::json similarityReport(const std::vector<SequenceFeature>& features) {
    nlohmann::json pairs = nlohmann::json::array();
    for (size_t i = 0; i < features.size(); ++i) {
        for (size_t j = i + 1; j < features.size(); ++j) {
            pairs.push_back({
                {"left", identity(features[i])},
                {"right", identity(features[j])},
                {"gait_similarity", dotProduct(features[i].gait, features[j].gait)},
                {"reid_similarity", dotProduct(features[i].reid, features[j].reid)},
                {"face_similarity", dotProduct(features[i].face, features[j].face)},
            });
        }
    }
    return {
        {"sequence_count", features.size()},
        {"pair_count", pairs.size()},
        {"pairs", pairs},
    };
}

static nlohmann::json parseSequenceDir(const fs::path& seqDir, const fs::path& outDir, std::vector<SequenceFeature>& features, int index) {
    std::vector<fs::path> frames = collectFrames(seqDir);
    if (frames.empty()) throw std::runtime_error("no frames under " + seqDir.string());

    nlohmann::json createPayload = {{"frame_count", frames.size()}};
    nlohmann::json created = nlohmann::json::parse(requestJSON("POST", "/v1/sequences", &createPayload).body);
    std::string taskID = created.at("task_id").get<std::string>();
    const nlohmann::json& upload = created.at("upload");
    if (upload.at("frame_count").get<size_t>() != frames.size()) throw std::runtime_error("upload frame count mismatch");
    uploadBatch(upload.at("upload_url").get<std::string>(), upload.at("upload_token").get<std::string>(), frames);
    nlohmann::json parsePayload = nlohmann::json::object();
    nlohmann::json parsed = nlohmann::json::parse(requestJSON("POST", "/v1/sequences/" + taskID + "/parse", &parsePayload).body);
    nlohmann::json record = {
        {"ok", true},
        {"sequence_dir", seqDir.string()},
        {"task_id", taskID},
        {"frame_count", frames.size()},
        {"result", parsed},
    };
    std::ostringstream name;
    name << "sequence_" << std::setw(3) << std::setfill('0') << index << ".json";
    fs::path outPath = outDir / name.str();
    writeJSON(outPath, record);

    for (const auto& seq : parsed.value("sequences", nlohmann::json::array())) {
        SequenceFeature feature;
        feature.source = outPath.string();
        feature.sequence_id = seq.value("sequence_id", "");
        feature.frame_count = seq.value("frame_count", 0);
        feature.gait = numericVector(seq.value("gait_feature", nlohmann::json::array()));
        feature.reid = numericVector(seq.value("reid_feature", nlohmann::json::array()));
        feature.face = numericVector(seq.value("face_feature", nlohmann::json::array()));
        features.push_back(std::move(feature));
    }
    return record;
}

} // namespace

int main(int argc, char** argv) {
    if (argc < 2 || argc > 3) {
        std::cerr << "usage: " << argv[0] << " sequence_root [output_dir]\n";
        return 2;
    }

    try {
        curl_global_init(CURL_GLOBAL_DEFAULT);
        fs::path sequenceRoot = argv[1];
        fs::path outputDir = argc >= 3 ? fs::path(argv[2]) : fs::path("output");
        fs::create_directories(outputDir);

        std::vector<fs::path> dirs = collectLeafSequenceDirs(sequenceRoot);
        if (dirs.empty()) throw std::runtime_error("no sequence dirs under " + sequenceRoot.string());

        nlohmann::json summary = {
            {"base_url", baseURL()},
            {"sequence_root", sequenceRoot.string()},
            {"sequence_count", dirs.size()},
            {"items", nlohmann::json::array()},
        };
        std::vector<SequenceFeature> features;
        int index = 0;
        for (const auto& dir : dirs) {
            ++index;
            std::cout << "[" << index << "/" << dirs.size() << "] " << dir << "\n";
            try {
                nlohmann::json item = parseSequenceDir(dir, outputDir, features, index);
                summary["items"].push_back({
                    {"ok", true},
                    {"sequence_dir", dir.string()},
                    {"task_id", item.value("task_id", "")},
                });
            } catch (const std::exception& ex) {
                summary["items"].push_back({
                    {"ok", false},
                    {"sequence_dir", dir.string()},
                    {"error", ex.what()},
                });
            }
            writeJSON(outputDir / "summary.json", summary);
        }
        writeJSON(outputDir / "similarity.json", similarityReport(features));
        curl_global_cleanup();
        return 0;
    } catch (const std::exception& ex) {
        curl_global_cleanup();
        std::cerr << ex.what() << "\n";
        return 1;
    }
}
