// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/flags/usage.h"
#include "absl/log/initialize.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "crypto/default/pem.h"
#include "google/protobuf/text_format.h"
#include "proto/generator.pb.h"
#include "testing/local_key_crypto_write_handler.h"
#include "tools/asset_writer.h"

ABSL_FLAG(std::string, input, "",
          "Path to the input asset file to read from (Required).");
ABSL_FLAG(std::string, output, "",
          "Path to the output asset file to write to (Required).");
ABSL_FLAG(std::string, claim_signer_key, "",
          "Path to PEM file containing the claim signer's EC P-256 private key "
          "(Required).");
ABSL_FLAG(std::string, claim_signer_cert, "",
          "Path to PEM file containing the claim signer's certificate chain, "
          "leaf first (Required).");
ABSL_FLAG(std::string, claim_signer_ocsp, "",
          "Path to DER file containing an OCSP response for the claim signer "
          "certificate (Optional).");
ABSL_FLAG(std::string, tsa_key, "",
          "Path to PEM file containing the TSA's EC P-256 private key "
          "(Optional; required unless timestamping is disabled).");
ABSL_FLAG(std::string, tsa_cert, "",
          "Path to PEM file containing the TSA's certificate "
          "(Optional; required unless timestamping is disabled).");
ABSL_FLAG(bool, disable_timestamping, false,
          "If true, disables timestamping when writing the C2PA manifest store "
          "(Default: false).");
ABSL_FLAG(std::string, request_txtpb, "",
          "Path to textproto file containing GenerateManifestStoreRequest "
          "(Optional; defaults to a basic 'created' action).");

namespace {

absl::StatusOr<std::string> GetContents(absl::string_view file_path) {
  std::ifstream f(std::string(file_path), std::ios::binary);
  if (!f.is_open()) {
    return absl::NotFoundError(
        absl::StrCat("Failed to open file: ", file_path));
  }
  std::stringstream buffer;
  buffer << f.rdbuf();
  return buffer.str();
}

}  // namespace

int main(int argc, char* argv[]) {
  absl::SetProgramUsageMessage(
      "Writes C2PA manifest stores into asset files.\n"
      "Usage:\n  asset_write --input=<input_path> --output=<output_path> "
      "--claim_signer_key=<pem_path> --claim_signer_cert=<pem_path> "
      "[--claim_signer_ocsp=<der_path>] [--tsa_key=<pem_path>] "
      "[--tsa_cert=<pem_path>] [--disable_timestamping] "
      "[--request_txtpb=<txtpb_path>]");
  absl::InitializeLog();
  absl::ParseCommandLine(argc, argv);

  const std::string input_path = absl::GetFlag(FLAGS_input);
  if (input_path.empty()) {
    std::cerr << "Error: --input flag is required.\n";
    return 1;
  }

  const std::string output_path = absl::GetFlag(FLAGS_output);
  if (output_path.empty()) {
    std::cerr << "Error: --output flag is required.\n";
    return 1;
  }

  const std::string claim_signer_key_path =
      absl::GetFlag(FLAGS_claim_signer_key);
  if (claim_signer_key_path.empty()) {
    std::cerr << "Error: --claim_signer_key flag is required.\n";
    return 1;
  }

  const std::string claim_signer_cert_path =
      absl::GetFlag(FLAGS_claim_signer_cert);
  if (claim_signer_cert_path.empty()) {
    std::cerr << "Error: --claim_signer_cert flag is required.\n";
    return 1;
  }

  credentio::LocalKeyCryptoWriteHandlerOptions crypto_options;

  absl::StatusOr<std::string> key_pem = GetContents(claim_signer_key_path);
  if (!key_pem.ok()) {
    std::cerr << "Error reading claim signer key file: " << key_pem.status()
              << "\n";
    return 1;
  }
  crypto_options.claim_signer.key_pem = *std::move(key_pem);

  absl::StatusOr<std::string> cert_pem = GetContents(claim_signer_cert_path);
  if (!cert_pem.ok()) {
    std::cerr << "Error reading claim signer cert file: " << cert_pem.status()
              << "\n";
    return 1;
  }
  absl::StatusOr<std::vector<std::string>> cert_chain =
      credentio::LoadCertsFromPem(*cert_pem);
  if (!cert_chain.ok()) {
    std::cerr << "Error parsing claim signer certificate chain: "
              << cert_chain.status() << "\n";
    return 1;
  }
  crypto_options.claim_signer.certificate_chain = *std::move(cert_chain);

  const std::string claim_signer_ocsp_path =
      absl::GetFlag(FLAGS_claim_signer_ocsp);
  if (!claim_signer_ocsp_path.empty()) {
    absl::StatusOr<std::string> ocsp_der = GetContents(claim_signer_ocsp_path);
    if (!ocsp_der.ok()) {
      std::cerr << "Error reading claim signer OCSP file: " << ocsp_der.status()
                << "\n";
      return 1;
    }
    crypto_options.claim_signer.ocsp_responses.push_back(*std::move(ocsp_der));
  }

  const std::string tsa_key_path = absl::GetFlag(FLAGS_tsa_key);
  const std::string tsa_cert_path = absl::GetFlag(FLAGS_tsa_cert);
  bool disable_timestamping = absl::GetFlag(FLAGS_disable_timestamping);

  if (!disable_timestamping && tsa_key_path.empty() && tsa_cert_path.empty()) {
    std::cerr << "WARNING: --tsa_key and --tsa_cert are not provided. "
                 "Disabling timestamping.\n";
    disable_timestamping = true;
  }

  crypto_options.disable_timestamping = disable_timestamping;
  if (!disable_timestamping) {
    if (tsa_key_path.empty() || tsa_cert_path.empty()) {
      std::cerr << "Error: Both --tsa_key and --tsa_cert must be provided when "
                   "timestamping is enabled.\n";
      return 1;
    }
    absl::StatusOr<std::string> tsa_key_pem = GetContents(tsa_key_path);
    if (!tsa_key_pem.ok()) {
      std::cerr << "Error reading TSA key file: " << tsa_key_pem.status()
                << "\n";
      return 1;
    }
    crypto_options.tsa.key_pem = *std::move(tsa_key_pem);

    absl::StatusOr<std::string> tsa_cert_pem = GetContents(tsa_cert_path);
    if (!tsa_cert_pem.ok()) {
      std::cerr << "Error reading TSA cert file: " << tsa_cert_pem.status()
                << "\n";
      return 1;
    }
    crypto_options.tsa.cert_pem = *std::move(tsa_cert_pem);
  }

  credentio::GenerateManifestStoreRequest request;
  const std::string request_txtpb_path = absl::GetFlag(FLAGS_request_txtpb);
  if (!request_txtpb_path.empty()) {
    absl::StatusOr<std::string> request_txtpb = GetContents(request_txtpb_path);
    if (!request_txtpb.ok()) {
      std::cerr << "Error reading request txtpb file: "
                << request_txtpb.status() << "\n";
      return 1;
    }
    if (!google::protobuf::TextFormat::ParseFromString(*request_txtpb,
                                                       &request)) {
      std::cerr << "Error parsing GenerateManifestStoreRequest from txtpb.\n";
      return 1;
    }
  } else {
    request.mutable_actions()->mutable_created();
  }

  absl::Status status = credentio::WriteAssetFile(
      std::move(crypto_options), std::move(request), input_path, output_path);
  if (!status.ok()) {
    std::cerr << "Asset write failed: " << status << "\n";
    return 1;
  }

  std::cout << "Successfully wrote C2PA asset to " << output_path << "\n";
  return 0;
}