#include <fstream>
#include <sstream>
#include <stdexcept>
#include <iomanip>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/sha.h> // OpenSSL for hashing
#include <openssl/evp.h> // OpenSSL for Base64 encoding

#include "ImageSigner.hpp"

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

ImageSigner::ImageSigner(const std::string &certsPath, const std::string &privateKeyPath)
    : certsPath(certsPath), privateKeyPath(privateKeyPath)
{

    originalImageSize = cv::Size(0, 0);
    thumbnailSize = cv::Size(0, 0);

    signInfo = std::make_unique<c2pa::SignerInfo>();
    signInfo->alg = "es256";
    signInfo->sign_cert = ReadTextFile(certsPath).c_str();
    signInfo->private_key = ReadTextFile(privateKeyPath).c_str();
    signInfo->ta_url = "http://timestamp.digicert.com";
}

void ImageSigner::GenerateThumbnailAndHash(const std::string &inputImagePath, int width, int height)
{
    cv::Mat image = cv::imread(inputImagePath);
    if (image.empty())
    {
        throw ThumbnailCreationException("Could not open or find the image at " + inputImagePath);
    }

    originalImageSize = image.size();
    thumbnailSize = cv::Size(width, height);

    // Generate the original hash and signed base64 hash
    try
    {
        std::vector<uchar> originalBuf;
        cv::imencode(".jpg", image, originalBuf);
        std::string originalData(originalBuf.begin(), originalBuf.end());
        std::string originalHash = GenerateSHA256Hash(originalData);
        originalImageBase64SignedHash = Base64Encode(SignHash(originalHash));
    }
    catch (const std::exception &e)
    {
        throw OriginalImageHashException(e.what());
    }

    // Create the thumbnail by resizing the image
    cv::Mat thumbnail;
    cv::resize(image, thumbnail, thumbnailSize);

    // Set the JPEG compression quality (e.g., 90)
    std::vector<int> compressionParams = {cv::IMWRITE_JPEG_QUALITY, 90};

    // Encode thumbnail to a buffer (in-memory) with specified quality
    std::vector<uchar> buf;
    cv::imencode(".jpg", thumbnail, buf, compressionParams);

    // Convert buffer to Base64 and generate the signed hash
    try
    {
        std::string thumbnailData(buf.begin(), buf.end());
        thumbnailBase64 = Base64Encode(thumbnailData);
        thumbnailBase64SignedHash = Base64Encode(SignHash(GenerateSHA256Hash(thumbnailData)));
    }
    catch (const std::exception &e)
    {
        throw Base64EncodingException(e.what());
    }
}

void ImageSigner::SignImage(const std::string &inputImagePath, const std::string &outputImagePath, nlohmann::json &claim)
{
    try
    {
        std::string claimJson = claim.dump();
        c2pa::sign_file(inputImagePath.c_str(), outputImagePath.c_str(), claimJson.c_str(), signInfo.get());
    }
    catch (const c2pa::Exception &e)
    {
        throw std::runtime_error("C2PA Signing Error: " + std::string(e.what()));
    }
}

cv::Size ImageSigner::GetThumbnailSize() const
{
    return thumbnailSize;
}

std::string ImageSigner::GetThumbnailBase64() const
{
    return thumbnailBase64;
}

std::string ImageSigner::GetThumbnailBase64SignedHash() const
{
    return thumbnailBase64SignedHash;
}

cv::Size ImageSigner::GetOriginalImageSize() const
{
    return originalImageSize;
}

std::string ImageSigner::GetOriginalImageBase64SignedHash() const
{
    return originalImageBase64SignedHash;
}

string ImageSigner::ReadTextFile(const std::string &path)
{
    ifstream file(path);
    if (!file.is_open())
    {
        throw FileReadException("Could not open file at " + path);
    }
    string contents((istreambuf_iterator<char>(file)), istreambuf_iterator<char>());
    file.close();
    return contents;
}

string ImageSigner::GenerateSHA256Hash(const std::string &data)
{
    unsigned char hash[SHA256_DIGEST_LENGTH];
    if (!SHA256(reinterpret_cast<const unsigned char *>(data.c_str()), data.size(), hash))
    {
        throw runtime_error("Failed to generate SHA-256 hash");
    }
    stringstream ss;
    for (int i = 0; i < SHA256_DIGEST_LENGTH; ++i)
    {
        ss << hex << setw(2) << setfill('0') << (int)hash[i];
    }
    return ss.str();
}

std::string ImageSigner::Base64Encode(const std::string &data)
{
    BIO *bio, *b64;
    BUF_MEM *bufferPtr;
    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);
    BIO_write(bio, data.c_str(), data.length());
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bufferPtr);
    BIO_set_close(bio, BIO_NOCLOSE);
    BIO_free_all(bio);
    return std::string(bufferPtr->data, bufferPtr->length);
}
std::string ImageSigner::SignHash(const std::string &hash)
{
    // Create a new EVP_PKEY object from the private key file
    EVP_PKEY *pkey = nullptr;
    FILE *pkey_file = fopen(privateKeyPath.c_str(), "r");
    if (!pkey_file)
    {
        throw std::runtime_error("Failed to open private key file at: " + privateKeyPath);
    }
    pkey = PEM_read_PrivateKey(pkey_file, nullptr, nullptr, nullptr);
    fclose(pkey_file);
    if (!pkey)
    {
        char errbuf[256];
        ERR_error_string_n(ERR_get_error(), errbuf, sizeof(errbuf));
        std::cerr << "Failed to read private key: " << errbuf << std::endl;
        throw std::runtime_error("Failed to read private key");
    }

    // Create a new EVP_MD_CTX to hold the signing context
    EVP_MD_CTX *mdctx = EVP_MD_CTX_new();
    if (!mdctx)
    {
        EVP_PKEY_free(pkey);
        throw std::runtime_error("Failed to create EVP_MD_CTX");
    }

    // Initialize the signing operation
    if (EVP_DigestSignInit(mdctx, nullptr, EVP_sha256(), nullptr, pkey) <= 0)
    {
        EVP_MD_CTX_free(mdctx);
        EVP_PKEY_free(pkey);
        throw std::runtime_error("Failed to initialize signing context");
    }

    // Perform the signing operation
    size_t sig_len = 0;
    if (EVP_DigestSign(mdctx, nullptr, &sig_len, reinterpret_cast<const unsigned char *>(hash.c_str()), hash.size()) <= 0)
    {
        EVP_MD_CTX_free(mdctx);
        EVP_PKEY_free(pkey);
        throw std::runtime_error("Failed to calculate signature length");
    }

    std::vector<unsigned char> signature(sig_len);
    if (EVP_DigestSign(mdctx, signature.data(), &sig_len, reinterpret_cast<const unsigned char *>(hash.c_str()), hash.size()) <= 0)
    {
        EVP_MD_CTX_free(mdctx);
        EVP_PKEY_free(pkey);
        throw std::runtime_error("Failed to sign the hash");
    }

    // Clean up
    EVP_MD_CTX_free(mdctx);
    EVP_PKEY_free(pkey);

    return std::string(signature.begin(), signature.end());
}

std::string ImageSigner::GetCurrentTimestamp() const
{
    std::time_t t = std::time(nullptr);
    std::tm tm;
#ifdef _WIN32
    gmtime_s(&tm, &t);
#else
    gmtime_r(&t, &tm);
#endif

    std::ostringstream oss;
    oss << std::put_time(&tm, "%Y-%m-%dT%H:%M:%SZ");
    return oss.str();
}