#include "hasher.hpp"
#include <mbedtls/md5.h>

namespace sphaira::hash {
namespace {

consteval auto CalculateHashStrLen(s64 buf_size) {
    return buf_size * 2 + 1;
}

struct FileSource final : BaseSource {
    FileSource(fs::Fs* fs, const fs::FsPath& path) : m_fs{fs} {
        m_open_result = m_fs->OpenFile(path, FsOpenMode_Read, std::addressof(m_file));
    }

    Result Size(s64* out) {
        return m_file.GetSize(out);
    }

    Result Read(void* buf, s64 off, s64 size, u64* bytes_read) {
        return m_file.Read(off, buf, size, 0, bytes_read);
    }

private:
    fs::Fs* m_fs{};
    fs::File m_file{};
    Result m_open_result{};
};

struct HashSource {
    virtual ~HashSource() = default;
    virtual void Update(const void* buf, s64 size) = 0;
    virtual void Get(std::string& out) = 0;
};

struct HashCrc32 final : HashSource {
    void Update(const void* buf, s64 size) {
        m_seed = crc32CalculateWithSeed(m_seed, buf, size);
    }

    void Get(std::string& out) {
        char str[CalculateHashStrLen(sizeof(m_seed))];
        std::snprintf(str, sizeof(str), "%08x", m_seed);
        out = str;
    }

private:
    u32 m_seed{};
};

struct HashMd5 final : HashSource {
    HashMd5() {
        mbedtls_md5_init(&m_ctx);
        mbedtls_md5_starts_ret(&m_ctx);
    }

    ~HashMd5() {
        mbedtls_md5_free(&m_ctx);
    }

    void Update(const void* buf, s64 size) {
        mbedtls_md5_update_ret(&m_ctx, (const u8*)buf, size);
    }

    void Get(std::string& out) {
        u8 hash[16];
        mbedtls_md5_finish_ret(&m_ctx, hash);

        char str[CalculateHashStrLen(sizeof(hash))];
        for (u32 i = 0; i < sizeof(hash); i++) {
            std::sprintf(str + i * 2, "%02x", hash[i]);
        }

        out = str;
    }

private:
    mbedtls_md5_context m_ctx{};
};

struct HashSha1 final : HashSource {
    HashSha1() {
        sha1ContextCreate(&m_ctx);
    }

    void Update(const void* buf, s64 size) {
        sha1ContextUpdate(&m_ctx, buf, size);
    }

    void Get(std::string& out) {
        u8 hash[SHA1_HASH_SIZE];
        sha1ContextGetHash(&m_ctx, hash);

        char str[CalculateHashStrLen(sizeof(hash))];
        for (u32 i = 0; i < sizeof(hash); i++) {
            std::sprintf(str + i * 2, "%02x", hash[i]);
        }

        out = str;
    }

private:
    Sha1Context m_ctx{};
};

struct HashSha256 final : HashSource {
    HashSha256() {
        sha256ContextCreate(&m_ctx);
    }

    void Update(const void* buf, s64 size) {
        sha256ContextUpdate(&m_ctx, buf, size);
    }

    void Get(std::string& out) {
        u8 hash[SHA256_HASH_SIZE];
        sha256ContextGetHash(&m_ctx, hash);

        char str[CalculateHashStrLen(sizeof(hash))];
        for (u32 i = 0; i < sizeof(hash); i++) {
            std::sprintf(str + i * 2, "%02x", hash[i]);
        }

        out = str;
    }

private:
    Sha256Context m_ctx{};
};

Result Hash(ui::ProgressBox* pbox, std::unique_ptr<HashSource> hash, std::shared_ptr<BaseSource> source, std::string& out) {
    s64 size;
    R_TRY(source->Size(&size));

    s64 offset{};
    std::vector<u8> chunk(1024 * 512);
    while (offset < size) {
        R_TRY(pbox->ShouldExitResult());
        const auto rsize = std::min<s64>(chunk.size(), size - offset);

        u64 bytes_read;
        R_TRY(source->Read(chunk.data(), offset, rsize, &bytes_read));
        hash->Update(chunk.data(), bytes_read);

        offset += bytes_read;
        pbox->UpdateTransfer(offset, size);
    }

    hash->Get(out);
    R_SUCCEED();
}

} // namespace

auto GetTypeStr(Type type) -> const char* {
    switch (type) {
        case Type::Crc32: return "CRC32";
        case Type::Md5: return "MD5";
        case Type::Sha1: return "SHA1";
        case Type::Sha256: return "SHA256";
    }
    return "";
}

Result Hash(ui::ProgressBox* pbox, Type type, std::shared_ptr<BaseSource> source, std::string& out) {
    switch (type) {
        case Type::Crc32: return Hash(pbox, std::make_unique<HashCrc32>(), source, out);
        case Type::Md5: return Hash(pbox, std::make_unique<HashMd5>(), source, out);
        case Type::Sha1: return Hash(pbox, std::make_unique<HashSha1>(), source, out);
        case Type::Sha256: return Hash(pbox, std::make_unique<HashSha256>(), source, out);
    }
    R_THROW(0x1);
}

Result Hash(ui::ProgressBox* pbox, Type type, fs::Fs* fs, const fs::FsPath& path, std::string& out) {
    auto source = std::make_shared<FileSource>(fs, path);
    return Hash(pbox, type, source, out);
}

} // namespace sphaira::has
