zenith-cli/main.cpp

#include <iostream>
#include <string>
#include <filesystem>
#include <cstdlib>
#include <vector>
#include <fstream>
#include <iomanip>
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <opencv2/opencv.hpp>
#include <curl/curl.h>

extern "C" {
    #include "libs/argparse.h"
}

namespace fs = std::filesystem;

struct ColorInfo {
    int r, g, b;
    double brightness;
    std::string hex;
};

double getColorDistance(const ColorInfo& c1, const ColorInfo& c2) {
    return std::sqrt(std::pow(c1.r - c2.r, 2) + std::pow(c1.g - c2.g, 2) + std::pow(c1.b - c2.b, 2));
}

void extract_color_palette(const std::string& image_path, int num_colors, const std::string& output_path) {
    if (num_colors < 3) {
        std::cerr << "ошибка: для генерации палитры нужно минимум 3 цвета" << std::endl;
        return;
    }

    cv::Mat image = cv::imread(image_path);
    if (image.empty()) {
        std::cerr << "ошибка: невозможно прочитать картинку по пути " << image_path << std::endl;
        return;
    }

    cv::Mat data;
    image.convertTo(data, CV_32F);
    
    data = data.reshape(3, data.total());

    cv::Mat labels;
    cv::Mat centers;
    cv::TermCriteria criteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 100, 1.0);
    
    cv::kmeans(data, num_colors, labels, criteria, 3, cv::KMEANS_PP_CENTERS, centers);

    std::vector<ColorInfo> colors;
    std::vector<ColorInfo> pastel_colors;

    for (int i = 0; i < centers.rows; i++) {
        ColorInfo c;
        c.b = static_cast<int>(centers.at<float>(i, 0));
        c.g = static_cast<int>(centers.at<float>(i, 1));
        c.r = static_cast<int>(centers.at<float>(i, 2));
        c.brightness = 0.299 * c.r + 0.587 * c.g + 0.114 * c.b;

        char hex_str[8];
        snprintf(hex_str, sizeof(hex_str), "#%02x%02x%02x", c.r, c.g, c.b);
        c.hex = hex_str;
        colors.push_back(c);

        cv::Mat bgr_pixel(1, 1, CV_8UC3, cv::Scalar(c.b, c.g, c.r));
        cv::Mat hsv_pixel;
        cv::cvtColor(bgr_pixel, hsv_pixel, cv::COLOR_BGR2HSV);

        cv::Vec3b& hsv = hsv_pixel.at<cv::Vec3b>(0, 0);
        
        hsv[1] = std::min(static_cast<int>(hsv[1]), 60); 
        if (hsv[1] < 30) hsv[1] = 40;
        hsv[2] = std::max(static_cast<int>(hsv[2]), 230); 

        cv::Mat pastel_bgr;
        cv::cvtColor(hsv_pixel, pastel_bgr, cv::COLOR_HSV2BGR);
        cv::Vec3b pbgr = pastel_bgr.at<cv::Vec3b>(0, 0);

        ColorInfo pc;
        pc.b = pbgr[0]; pc.g = pbgr[1]; pc.r = pbgr[2];
        pc.brightness = 0.299 * pc.r + 0.587 * pc.g + 0.114 * pc.b;

        char hex_str_pas[8];
        snprintf(hex_str_pas, sizeof(hex_str_pas), "#%02x%02x%02x", pc.r, pc.g, pc.b);
        pc.hex = hex_str_pas;
        pastel_colors.push_back(pc);
    }

    std::sort(colors.begin(), colors.end(), [](const ColorInfo& a, const ColorInfo& b) {
        return a.brightness < b.brightness;
    });

    std::sort(pastel_colors.begin(), pastel_colors.end(), [](const ColorInfo& a, const ColorInfo& b) {
        return a.brightness < b.brightness;
    });

    ColorInfo bg_color = colors[0];
    ColorInfo bg_sub = colors[1];
    ColorInfo fg_color = colors[num_colors - 1];

    ColorInfo accent_color = colors[1]; 
    double max_distance = 0;
    for (int i = 1; i < num_colors - 1; i++) {
        double dist = getColorDistance(colors[i], bg_color);
        if (dist > max_distance) {
            max_distance = dist;
            accent_color = colors[i];
        }
    }

    ColorInfo fg_sub_color = colors[num_colors - 2]; 
    for (int i = num_colors - 2; i >= 1; i--) {
        if (colors[i].hex != accent_color.hex) {
            fg_sub_color = colors[i];
            break;
        }
    }

    ColorInfo pas_bg_color = pastel_colors[0];
    ColorInfo pas_bg_sub = pastel_colors[1];
    ColorInfo pas_fg_color = pastel_colors[num_colors - 1];
    
    ColorInfo pas_accent_color = pastel_colors[1];
    double max_dist_pas = 0;
    for (int i = 1; i < num_colors - 1; i++) {
        double dist = getColorDistance(pastel_colors[i], pas_bg_color);
        if (dist > max_dist_pas) {
            max_dist_pas = dist;
            pas_accent_color = pastel_colors[i];
        }
    }
    ColorInfo pas_fg_sub_color = pastel_colors[num_colors - 2];
    for (int i = num_colors - 2; i >= 1; i--) {
        if (pastel_colors[i].hex != pas_accent_color.hex) {
            pas_fg_sub_color = pastel_colors[i];
            break;
        }
    }

    fs::path out_dir(output_path);

    auto write_palette = [&](const std::string& filename, const std::vector<ColorInfo>& palette) {
        std::ofstream file(out_dir / filename);
        if (file.is_open()) {
            for (const auto& color : palette) {
                file << color.hex << "\n";
            }
            file.close();
            std::cout << "успешно сформирована палитра: " << filename << std::endl;
        } else {
            std::cerr << "ошибка: не удалось создать файл " << filename << std::endl;
        }
    };

    write_palette("palette_dark", {bg_color, bg_sub, accent_color, fg_sub_color, fg_color});
    write_palette("palette_light", {fg_color, fg_sub_color, accent_color, bg_sub, bg_color});
    write_palette("palette_pastel_dark", {pas_bg_color, pas_bg_sub, pas_accent_color, pas_fg_sub_color, pas_fg_color});
    write_palette("palette_pastel_light", {pas_fg_color, pas_fg_sub_color, pas_accent_color, pas_bg_sub, pas_bg_color});
}

int download_wall(const char* link, const char* path) {
    CURL* curl;
    FILE* fp;
    CURLcode res;

    curl_global_init(CURL_GLOBAL_ALL);
    curl = curl_easy_init();
    
    if(curl) {
        fp = fopen(path, "wb");
        if(fp == NULL) {
            fprintf(stderr, "ошибка: не удалось создать или открыть файл %s\n", path);
            curl_easy_cleanup(curl);
            curl_global_cleanup();
            return 1;
        }

        curl_easy_setopt(curl, CURLOPT_URL, link);
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, fp);
        curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
        
        curl_easy_setopt(curl, CURLOPT_USERAGENT, "Mozilla/5.0 (Windows NT 10.0; Win64; x64) WallpaperPaletteBot/1.0");

        res = curl_easy_perform(curl);

        fclose(fp);
        curl_easy_cleanup(curl);

        if(res != CURLE_OK) {
            fprintf(stderr, "ошибка скачивания: %s\n", curl_easy_strerror(res));
            curl_global_cleanup();
            return 1;
        }
    } else {
        curl_global_cleanup();
        return 1;
    }

    curl_global_cleanup();
    return 0;
}

int main(int argc, const char **argv) {
    const char* link = nullptr;
    const char* raw_scheme_path = nullptr;
    const char* raw_wall_path = nullptr;
    const char* raw_path_to_local = nullptr;

    struct argparse_option options[] = {
        OPT_HELP(),
        OPT_STRING(0, "link", &link, "ссылка на изображение с сайта wallhave.cc", NULL, 0, 0),
        OPT_STRING(0, "local", &raw_path_to_local, "путь до локального файла с обоями", NULL, 0, 0),
        OPT_STRING(0, "scheme-output", &raw_scheme_path, "директория для сохранения палитр", NULL, 0, 0),
        OPT_STRING(0, "wall-output", &raw_wall_path, "директория для сохранения обоев", NULL, 0, 0),
        OPT_END(),
    };

    struct argparse argparse;
    argparse_init(&argparse, options, nullptr, 0);
    argc = argparse_parse(&argparse, argc, argv);

    if (link == nullptr && raw_path_to_local == nullptr) {
        std::cerr << "ошибка: необходимо указать либо --link, либо --local" << std::endl;
        return 1;
    }
    if (link != nullptr && raw_path_to_local != nullptr) {
        std::cerr << "ошибка: нельзя указать одновременно --link и --local. Выберите что-то одно." << std::endl;
        return 1;
    }

    auto resolve_dir = [](const char* raw_path) -> fs::path {
        if (raw_path == nullptr || std::string(raw_path).empty()) {
            return fs::current_path();
        }
        std::string path_str(raw_path);
        if (path_str[0] == '~') {
            const char* user_home = std::getenv("HOME");
            if (user_home) path_str.replace(0, 1, user_home);
        }
        return fs::path(path_str);
    };

    fs::path target_dir_wall = resolve_dir(raw_wall_path);
    if (!fs::exists(target_dir_wall)) {
        fs::create_directories(target_dir_wall);
    }

    fs::path target_dir_scheme = resolve_dir(raw_scheme_path);
    if (!fs::exists(target_dir_scheme)) {
        fs::create_directories(target_dir_scheme);
    }

    std::string final_image_path;

    if (raw_path_to_local != nullptr) {
        std::string local_str(raw_path_to_local);
        
        if (local_str[0] == '~') {
            const char* user_home = std::getenv("HOME");
            if (user_home) local_str.replace(0, 1, user_home);
        }
        
        fs::path local_file(local_str);
        
        if (!fs::exists(local_file)) {
            std::cerr << "ошибка: локальный файл не найден по пути " << local_file << std::endl;
            return 1;
        }
        
        final_image_path = local_file.string();
        std::cout << "используется локальный файл: " << final_image_path << std::endl;

    } else {
        std::string full_wall_path = (target_dir_wall / "wallpaper.jpg").string();
        
        std::cout << "скачивание изображения..." << std::endl;
        if (download_wall(link, full_wall_path.c_str()) != 0) {
            std::cerr << "не удалось скачать изображение. завершение работы." << std::endl;
            return 1;
        }
        
        final_image_path = full_wall_path;
    }

    std::cout << "извлечение палитры..." << std::endl;
    extract_color_palette(final_image_path, 6, target_dir_scheme.string()); 

    return 0;
}