// embed_faces — run SCRFD-500MF + ArcFace on a list of image files and write
// embeddings as JSON to stdout.
//
// Usage:
//   embed_faces [--detector <path>] [--arcface <path>] [--conf <f>] [--nms <f>]
//               image1.jpg image2.jpg ...
//
// Output (stdout): JSON array, one object per input image:
//   [
//     {
//       "image":      "actor.jpg",
//       "embedding":  [0.012, -0.034, ...],   // 512 floats, L2-normalised
//       "bbox":       [x, y, w, h],
//       "confidence": 0.91
//     },
//     {
//       "image":     "bad.jpg",
//       "embedding": null,                    // no face detected / alignment failed
//       "error":     "no face detected"
//     }
//   ]
//
// Design: each image is processed independently.  If multiple faces are
// detected the one with the highest confidence is used (gallery images are
// expected to contain exactly one subject).  A warning is printed to stderr
// when more than one face is found.
//
// --all-faces emits every detection instead, which is what a caller analysing
// a frame rather than a gallery portrait needs:
//   [ { "image": "frame.png",
//       "faces": [ {"bbox": [...], "landmarks": [[x,y] x5],
//                   "confidence": 0.89, "embedding": [...]} , ... ] } ]
//
// --calibration <gallery> additionally emits the gallery's fitted Platt
// sigmoid, so a non-Python client can turn a similarity into P(match) with the
// same parameters the C++ matcher uses.  Output becomes
// {"calibration": {...}, "images": [...]}.  Clients must score through it:
// AR-024 requires the calibrated probability, never a bare cosine — a raw
// threshold means something different for every model, gallery and face size.
//
// This binary is intentionally a thin wrapper around the same ONNX models
// used by scene_analyze, so embeddings are guaranteed compatible.

#include "config.hpp"
#include "face_utils.hpp"
#include "gallery/gallery_store.hpp"
#include "inference/face_detector.hpp"
#include "inference/face_embedder.hpp"

#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

#include <nlohmann/json.hpp>

#include <algorithm>
#include <cstring>
#include <filesystem>
#include <functional>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

namespace fs = std::filesystem;

using json = nlohmann::json;

// ── Per-image result ──────────────────────────────────────────────────────────

struct FaceResult {
    std::string image_path;
    bool        ok{false};
    std::string error;
    Embedding   embedding{};
    float       confidence{0.f};
    float       bbox[4]{};   // x, y, w, h
    std::array<cv::Point2f, 5> landmarks{};
};

// ── Debug rendering ───────────────────────────────────────────────────────────
// Writes <dir>/<stem>_annotated.jpg (input with bbox + 5 landmarks) and
// <dir>/<stem>_aligned.jpg (112×112 aligned crop). Stem is derived from the
// parent directory and filename so images from different actor folders don't
// collide when fed into a single debug dir.

static std::string debug_stem(const std::string& path) {
    fs::path p(path);
    std::string parent = p.parent_path().filename().string();
    std::string stem   = p.stem().string();
    return parent.empty() ? stem : parent + "_" + stem;
}

static void save_debug(const std::string& dir,
                        const std::string& src_path,
                        const cv::Mat&     img,
                        const DetectedFace& face,
                        const cv::Mat&     aligned) {
    fs::create_directories(dir);
    cv::Mat annotated = img.clone();
    cv::rectangle(annotated, face.bbox, {0, 255, 0}, 2);
    static const cv::Scalar colors[5] = {
        {  0,   0, 255},   // right eye   — red
        {255,   0,   0},   // left eye    — blue
        {  0, 255, 255},   // nose        — yellow
        {  0, 255,   0},   // right mouth — green
        {255,   0, 255},   // left mouth  — magenta
    };
    for (int i = 0; i < 5; ++i)
        cv::circle(annotated, face.landmarks[i], 4, colors[i], -1);

    const std::string stem = debug_stem(src_path);
    cv::imwrite(dir + "/" + stem + "_annotated.jpg", annotated);
    cv::imwrite(dir + "/" + stem + "_aligned.jpg",   aligned);
}

// ── Process one image, keeping every detection ────────────────────────────────
// The --all-faces path. Same detect → align → embed chain as process() below,
// but without the highest-confidence reduction: a frame legitimately contains
// several people, and dropping all but one is a gallery-portrait assumption.
// Faces that fail alignment are reported with a null embedding rather than
// silently dropped, so a caller can count what the detector found against what
// survived the ArcFace warp.

struct MultiFaceResult {
    std::string             image_path;
    std::string             error;      // set only when the image itself failed
    std::vector<FaceResult> faces;
};

static MultiFaceResult process_all(
    const std::string& path,
    const std::function<std::vector<DetectedFace>(const cv::Mat&)>& detect,
    const std::function<Embedding(const cv::Mat&)>& embed_one,
    int                max_side,
    const std::string& debug_dir = "") {
    MultiFaceResult out;
    out.image_path = path;

    cv::Mat img = cv::imread(path);
    if (img.empty()) {
        out.error = "cannot read image";
        return out;
    }

    if (max_side > 0) {
        const int big = std::max(img.cols, img.rows);
        if (big > max_side) {
            const double s = static_cast<double>(max_side) / big;
            cv::resize(img, img, {}, s, s, cv::INTER_AREA);
        }
    }

    std::vector<DetectedFace> faces = detect(img);
    if (faces.empty()) {
        cv::Mat enhanced = enhance_for_retry(img);
        faces = detect(enhanced);
        if (!faces.empty())
            img = enhanced;
    }
    if (faces.empty()) {
        out.error = "no face detected";
        return out;
    }

    for (const auto& face : faces) {
        FaceResult r;
        r.image_path = path;
        r.confidence = face.confidence;
        r.bbox[0] = face.bbox.x;      r.bbox[1] = face.bbox.y;
        r.bbox[2] = face.bbox.width;  r.bbox[3] = face.bbox.height;
        r.landmarks = face.landmarks;

        cv::Mat crop = align_face(img, face.landmarks);
        if (crop.empty()) {
            r.error = "alignment failed";
        } else {
            r.ok        = true;
            r.embedding = embed_one(crop);
            if (!debug_dir.empty())
                save_debug(debug_dir, path, img, face, crop);
        }
        out.faces.push_back(std::move(r));
    }
    return out;
}

// ── Process one image ─────────────────────────────────────────────────────────

static FaceResult process(const std::string& path,
                           const std::function<std::vector<DetectedFace>(const cv::Mat&)>& detect,
                           const std::function<Embedding(const cv::Mat&)>& embed_one,
                           int                max_side,
                           const std::string& debug_dir = "") {
    FaceResult res;
    res.image_path = path;

    cv::Mat img = cv::imread(path);
    if (img.empty()) {
        res.error = "cannot read image";
        return res;
    }

    if (max_side > 0) {
        const int big = std::max(img.cols, img.rows);
        if (big > max_side) {
            const double s = static_cast<double>(max_side) / big;
            cv::resize(img, img, {}, s, s, cv::INTER_AREA);
        }
    }

    std::vector<DetectedFace> faces = detect(img);

    if (faces.empty()) {
        cv::Mat enhanced = enhance_for_retry(img);
        faces = detect(enhanced);
        if (!faces.empty())
            img = enhanced;
    }

    if (faces.empty()) {
        res.error = "no face detected";
        return res;
    }

    if (faces.size() > 1)
        std::cerr << "[warn] " << path << ": " << faces.size()
                  << " faces detected, using highest-confidence one\n";

    const auto& best = *std::max_element(
        faces.begin(), faces.end(),
        [](const DetectedFace& a, const DetectedFace& b) {
            return a.confidence < b.confidence;
        });

    cv::Mat crop = align_face(img, best.landmarks);
    if (crop.empty()) {
        res.error = "alignment failed";
        return res;
    }

    res.ok         = true;
    res.embedding  = embed_one(crop);
    res.confidence = best.confidence;
    res.bbox[0]    = best.bbox.x;
    res.bbox[1]    = best.bbox.y;
    res.bbox[2]    = best.bbox.width;
    res.bbox[3]    = best.bbox.height;
    res.landmarks  = best.landmarks;

    if (!debug_dir.empty())
        save_debug(debug_dir, path, img, best, crop);

    return res;
}

// ── Main ──────────────────────────────────────────────────────────────────────

int main(int argc, char** argv) {
    std::string detector_model = kDefaultDetectorModel;
    std::string detector_engine;
    std::string arcface_model  = kDefaultArcfaceModel;
    std::string arcface_engine;
    std::string debug_dir;
    std::string calibration_gallery;
    bool  all_faces = false;
    float conf = 0.5f, nms = 0.4f;
    int   max_side = 500;
    std::vector<std::string> images;

    for (int i = 1; i < argc; ++i) {
        if      (std::strcmp(argv[i], "--detector")         == 0 && i+1 < argc) { detector_model  = argv[++i]; }
        else if (std::strcmp(argv[i], "--detector-engine")  == 0 && i+1 < argc) { detector_engine = argv[++i]; }
        else if (std::strcmp(argv[i], "--arcface")        == 0 && i+1 < argc) { arcface_model  = argv[++i]; }
        else if (std::strcmp(argv[i], "--arcface-engine") == 0 && i+1 < argc) { arcface_engine = argv[++i]; }
        else if (std::strcmp(argv[i], "--conf")       == 0 && i+1 < argc) { conf = std::stof(argv[++i]); }
        else if (std::strcmp(argv[i], "--nms")        == 0 && i+1 < argc) { nms  = std::stof(argv[++i]); }
        else if (std::strcmp(argv[i], "--save-debug") == 0 && i+1 < argc) { debug_dir = argv[++i]; }
        else if (std::strcmp(argv[i], "--max-side")   == 0 && i+1 < argc) { max_side = std::stoi(argv[++i]); }
        else if (std::strcmp(argv[i], "--calibration")== 0 && i+1 < argc) { calibration_gallery = argv[++i]; }
        else if (std::strcmp(argv[i], "--all-faces")  == 0) { all_faces = true; }
        else if (argv[i][0] != '-') { images.push_back(argv[i]); }
        else { std::cerr << "[warn] unknown flag: " << argv[i] << "\n"; }
    }

    if (images.empty()) {
        std::cerr << "Usage: embed_faces [--detector <path>] [--arcface <path>] "
                     "[--save-debug <dir>] [--max-side <N>] [--all-faces] "
                     "[--calibration <gallery>] image1.jpg ...\n";
        return 1;
    }

    Config cfg;
    cfg.detector_model  = detector_model;
    cfg.detector_engine = detector_engine;
    cfg.arcface_model   = arcface_model;
    cfg.arcface_engine  = arcface_engine;
    cfg.detector_conf   = conf;
    cfg.detector_nms    = nms;

    // The compiled-in inference backend (ORT or TRT) is chosen by the factories.
    auto detector = make_face_detector(cfg);
    auto embedder = make_face_embedder(cfg);

    std::function<std::vector<DetectedFace>(const cv::Mat&)> detect =
        [&](const cv::Mat& im) { return detector->detect(im); };
    std::function<Embedding(const cv::Mat&)> embed_one =
        [&](const cv::Mat& c) { return embedder->embed_one(c); };

    // One face's fields, shared by both output shapes.
    auto face_json = [](const FaceResult& r) {
        json f;
        f["confidence"] = r.confidence;
        f["bbox"]       = {r.bbox[0], r.bbox[1], r.bbox[2], r.bbox[3]};
        json lms = json::array();
        for (const auto& pt : r.landmarks) lms.push_back({pt.x, pt.y});
        f["landmarks"] = std::move(lms);
        if (r.ok) f["embedding"] = std::vector<float>(r.embedding.begin(),
                                                      r.embedding.end());
        else    { f["embedding"] = nullptr; f["error"] = r.error; }
        return f;
    };

    // Process images and build JSON output
    json images_out = json::array();

    for (const auto& path : images) {
        std::cerr << "[embed_faces] " << path << "\n";

        json entry;
        entry["image"] = path;

        if (all_faces) {
            MultiFaceResult res = process_all(path, detect, embed_one, max_side, debug_dir);
            if (!res.error.empty()) {
                entry["faces"] = json::array();
                entry["error"] = res.error;
                std::cerr << "  [skip] " << res.error << "\n";
            } else {
                json faces = json::array();
                for (const auto& f : res.faces) faces.push_back(face_json(f));
                entry["faces"] = std::move(faces);
            }
        } else {
            FaceResult res = process(path, detect, embed_one, max_side, debug_dir);
            if (res.ok) {
                entry.merge_patch(face_json(res));
            } else {
                entry["embedding"] = nullptr;
                entry["error"]     = res.error;
                std::cerr << "  [skip] " << res.error << "\n";
            }
        }
        images_out.push_back(std::move(entry));
    }

    // Without --calibration the output stays a bare array, unchanged, so
    // existing callers (build_gallery, fetch_missing_actors) are unaffected.
    if (calibration_gallery.empty()) {
        std::cout << images_out.dump() << "\n";
        return 0;
    }

    ActorGallery gallery = load_gallery(calibration_gallery);
    if (!gallery.calib_valid)
        std::cerr << "[warn] " << calibration_gallery
                  << " carries no valid calibration; a client cannot convert a "
                     "similarity to a probability from it (AR-024)\n";

    json out;
    out["calibration"] = {
        {"a",     gallery.calib_a},
        {"b",     gallery.calib_b},
        {"valid", gallery.calib_valid},
        {"form",  "P(match) = 1/(1+exp(-(a*similarity + b + log_prior_odds)))"},
        {"note",  "Score through this. AR-024: a bare cosine threshold means "
                  "something different for every model, gallery and face size. "
                  "log_prior_odds = log(p0/(1-p0)) for a base-rate prior p0; use "
                  "0 for association (are these two faces one person)."},
    };
    out["images"] = std::move(images_out);
    std::cout << out.dump() << "\n";
    return 0;
}