#pragma once

#include <cstddef>
#include <cstdint>
#include <string>

#include "sha256.h"

namespace roro {

// The Update File (see CONTEXT.md): a 160-byte header, then the firmware image. All integers are
// little-endian.
//   0  magic "RORO-OTA"     8  u16 format (1)    10 u16 header size (160)   12 u32 image size
//  16  image SHA-256[32]   48  version, NUL-padded [32]
//  80  u16 signature length   82  signature (DER, up to 72 bytes)   154 reserved
// The signature covers SHA-256 of bytes 0..79, which include the image's hash, so a bad signature
// is caught before anything is written, and a bad image when its hash is checked at the end.
namespace update {
constexpr char kMagic[] = "RORO-OTA";
constexpr uint16_t kFormat = 1;
constexpr size_t kHeaderSize = 160;
constexpr size_t kSignedBytes = 80;
constexpr size_t kMaxSignature = 72;
}  // namespace update

class SignatureVerifier {
   public:
    virtual ~SignatureVerifier() = default;
    virtual bool verify(const uint8_t digest[32], const uint8_t* signature, size_t len) = 0;
};

// Where the image goes (the inactive app slot on the device).
class UpdateSink {
   public:
    virtual ~UpdateSink() = default;
    virtual bool begin(size_t imageSize) = 0;
    virtual bool write(const uint8_t* data, size_t len) = 0;
    virtual bool finish() = 0;  // make it the image to boot next
    virtual void abort() = 0;
};

// Streams an Update File into a sink: checks the header and signature first, then hashes the image
// as it passes through, and only finishes the sink if everything matches.
class UpdateParser {
   public:
    enum class State { Header, Image, Done, Failed };

    UpdateParser(SignatureVerifier& verifier, UpdateSink& sink, size_t maxImageSize, std::string installedVersion)
        : verifier_(verifier), sink_(sink), maxImage_(maxImageSize), installed_(std::move(installedVersion)) {}

    void feed(const uint8_t* data, size_t len);
    bool end();  // no more data: true if the update was installed
    // The whole image the header announced has arrived (the sender need not close the connection).
    bool complete() const { return state_ == State::Image && received_ == imageSize_; }

    State state() const { return state_; }
    const std::string& error() const { return error_; }
    const std::string& version() const { return version_; }
    bool isDowngrade() const { return downgrade_; }
    int percent() const { return imageSize_ ? static_cast<int>(received_ * 100 / imageSize_) : 0; }

   private:
    void parseHeader();
    void fail(const std::string& why);

    SignatureVerifier& verifier_;
    UpdateSink& sink_;
    size_t maxImage_;
    std::string installed_;

    State state_ = State::Header;
    uint8_t header_[update::kHeaderSize];
    size_t headerUsed_ = 0;
    uint8_t expectedHash_[32];
    size_t imageSize_ = 0;
    size_t received_ = 0;
    Sha256 hash_;
    std::string version_, error_;
    bool downgrade_ = false;
};

}  // namespace roro