#include <stdio.h>
#include <math.h>
#include "../cshared/copts.h"
#include "../cshared/cstr.h"
#include "ecc_api.h"
#include <ctype.h>

#include "asn_application.h"
#include "Etsits103097Certificate.h"
#include "xer_support.h"

#define CERT_MAX_SIZE 0x10000

static const char * _outPath = ".";
static const char * _searchPath = NULL;
static const char * _certName = NULL;
static       char * _profileName = NULL;
static       char * _signerName = NULL;
static ecc_format   _outKeyFormat = ecc_bin;
static const char * _verificationKey = NULL;
static const char * _decriptionKey   = NULL;
static const char * _keyPath = NULL;
static int          _force = 0;
static const char * _cfgFile = NULL;


EtsiTs103097Certificate_t * _cert = NULL;
EtsiTs103097Certificate_t * _issuer = NULL;

char   _tbsHash[256]; // has space for issuer hash
int    _tbsHashLength = 0;

char   _signerHash[256]; // has space for issuer hash
int    _signerHashLength = 0;

static const char * const _key_formats[] = {
	"bin", "hex", "pem", NULL
};

typedef struct {
	const char * cert;
	const char * vkey;
	const char * vkey_pub;
	const char * ekey;
	const char * ekey_pub;
} extensions_t;

#define EXT_CERT ".oer"
#define EXT_VKEY ".vkey"
#define EXT_EKEY ".ekey"
#define EXT_PUB  "_pub"

static copt_t _options [] = {
	{ "h?", "help", COPT_HELP, NULL,                 "Print this help page" },
	{ "C", "config", COPT_CFGFILE, (void*)&_cfgFile, "Config file path [no cfg file]" },
	{ "o", "out", COPT_STR, (void*)&_outPath,         "Output path [current dir by default]" },
	{ "k", "key-format", COPT_STRENUM, (void*)_key_formats, "Keys output format (bin|hex|pem)[binary by default]" },
	{ "S", "certs", COPT_STR, (void*)&_searchPath, "Certificates search path [Output path by default]" },
	{ "K", "keys",  COPT_STR,  (void*)&_keyPath,       "Private key storage path [Output path by default]" },
	{ "f", "force", COPT_BOOL, (void*)&_force,         "Force regenerate existing certificate and keys" },
	{ "n", "name", COPT_STR, (void*)&_certName,        "Certificate name (take from profile by default)" },
	{ "v", "vkey", COPT_STR, (void*)&_verificationKey, "Verification public key (generate key pair by default)" },
	{ "e", "ekey", COPT_STR, (void*)&_decriptionKey,   "Encription public key (generate key pair if neccessary)" },
	{ "s", "signer", COPT_STR, (void*)&_signerName,   "Signer certificate name [take from profile by default]" },
	{ NULL, NULL, COPT_END, NULL, NULL }
};

static int is_P384CurvePoint_empty(EccP384CurvePoint_t* point);
static int is_P256CurvePoint_empty(EccP256CurvePoint_t* point);

static void fill_curve_point_eccP256(EccP256CurvePoint_t* point, ecc_curve_id curveType, char * keyPath);
static void fill_curve_point_eccP384(EccP384CurvePoint_t* point, ecc_curve_id curveType, char * keyPath);

static int _issuer_parser_cb(pxml_chunk_type_e _type,
	const void *_chunk_data, size_t _chunk_size, void *_key)
{
	char * f = cstrnstr((const char *)_chunk_data, _chunk_size, "name=\"");
	if (f){
		_signerName = f + 6;
		f = cstrchr(_signerName, '"');
		_signerName = cstrndup(_signerName, f - _signerName);
	}
	return -1;	
}

static asn_dec_rval_t IssuerIdentifier_xer_decoder(const asn_codec_ctx_t *opt_codec_ctx,
	const asn_TYPE_descriptor_t *td, void **struct_ptr,
	const char *opt_mname, const void *buf_ptr, size_t size)
{
	if (_signerName == NULL){
		int stateContext = 0;
		pxml_parse(&stateContext, buf_ptr, size, _issuer_parser_cb, NULL);
	}
	return CHOICE_decode_xer(opt_codec_ctx, td, struct_ptr, opt_mname, buf_ptr, size);
}

typedef struct enc_to_buf_arg {
	void *buffer;
	size_t left;
} enc_to_buf_arg;

static asn_enc_rval_t ToBeSignedCertificate_oer_encoder(const asn_TYPE_descriptor_t *td,
	const asn_oer_constraints_t *constraints, const void *sptr,
	asn_app_consume_bytes_f *cb, void *app_key)
{
	asn_enc_rval_t rc;
	ToBeSignedCertificate_t * tbs = (ToBeSignedCertificate_t*)sptr;
	enc_to_buf_arg *a = (enc_to_buf_arg*)app_key;
	rc = SEQUENCE_encode_oer(td, constraints, sptr, cb, app_key);
	if (rc.encoded > 0){
		// calculate hash
		if (tbs->verifyKeyIndicator.present == VerificationKeyIndicator_PR_verificationKey &&
			tbs->verifyKeyIndicator.choice.verificationKey.present == PublicVerificationKey_PR_ecdsaBrainpoolP384r1) {
			_tbsHashLength = 48;
			sha384_calculate(_tbsHash, (const char*)a->buffer, rc.encoded);
		}
		else{
			_tbsHashLength = 32;
			sha256_calculate(_tbsHash, (const char*)a->buffer, rc.encoded);
		}
	}
	return rc;
}

static ecc_curve_id _pk_type_to_curveid[] = {
	0,
	ecies_nistp256,       //Signature_PR_ecdsaNistP256Signature,
	ecies_brainpoolp256r, //Signature_PR_ecdsaBrainpoolP256r1Signature,
	ecies_brainpoolp384r  //Signature_PR_ecdsaBrainpoolP384r1Signature
};

static ecc_curve_id _pk_type_to_hashid[] = {
	0,
	sha_256,       //Signature_PR_ecdsaNistP256Signature,
	sha_256, //Signature_PR_ecdsaBrainpoolP256r1Signature,
	sha_384  //Signature_PR_ecdsaBrainpoolP384r1Signature
};

static asn_enc_rval_t Signature_oer_encoder(const asn_TYPE_descriptor_t *td,
	const asn_oer_constraints_t *constraints, const void *sptr,
	asn_app_consume_bytes_f *cb, void *app_key)
{
	Signature_t * s = (Signature_t *)sptr;
	if (is_P256CurvePoint_empty(&s->choice.ecdsaNistP256Signature.rSig)){
		// look for signer private key
		ecc_curve_id alg = _pk_type_to_curveid[s->present];
		ecc_hash_id hashId = _pk_type_to_hashid[s->present];
		const char * sName = _signerName;

		if (sName == NULL && _cert->issuer.present == IssuerIdentifier_PR_self)
			sName = _profileName;

		char* pk = cvstrdup(_keyPath, "/", sName, EXT_VKEY, NULL);
		void *k = ecc_key_private_load(pk, alg);
		if (k){
			char h[48];
			int hl = 0;
			// calculate joint hash
			memcpy(_tbsHash, _signerHash, _signerHashLength);
			switch (hashId)
			{
			case sha_256:
				sha256_calculate(h, _tbsHash, _tbsHashLength + _signerHashLength);
				hl = sha256_hash_size;
				break;
			case sha_384:
				sha384_calculate(h, _tbsHash, _tbsHashLength + _signerHashLength);
				hl = sha384_hash_size;
				break;
			}
			OCTET_STRING_fromBuf(&s->choice.ecdsaNistP256Signature.rSig.choice.x_only, h, hl);
			OCTET_STRING_fromBuf(&s->choice.ecdsaNistP256Signature.sSig, h, hl);

			ecc_sign(k, h, hl, s->choice.ecdsaNistP256Signature.rSig.choice.x_only.buf, s->choice.ecdsaNistP256Signature.sSig.buf);
		}
	}
	return asn_OP_CHOICE.oer_encoder(td, constraints, sptr, cb, app_key);
}

int main(int argc, char ** argv)
{
	// set default time to the begining of this year
	//_setup_default_time();

	//parse options
	argc = coptions(argc, argv, COPT_HELP_NOVALUES , _options);
	if(argc < 2){
		if(argc<0 && (0-argc)<((sizeof(_options)/sizeof(_options[0]))-1)){
			printf("Unknown option %s\n", argv[0-argc]);
		}
		const char * a = strrchr(argv[0], '/');
		if (a == NULL) a = argv[0];
		coptions_help(stdout, a, COPT_HELP_NOVALUES, _options, "<profile> [signer]");
		return -1;
	}

	if (_searchPath == NULL) _searchPath = _outPath;
	if (_keyPath == NULL)    _keyPath    = _outPath;
	
	_outKeyFormat = copts_enum_value(_options, 3, _key_formats);
	
	if(argc > 2){
		// set signer certificate file name
		_signerName = argv[2];
	}
	
	if(ecc_api_init()){
		return -1;
	}

	_profileName = cstrdup(cstrlastpathelement(argv[1]));
	if(_profileName){
		char * p = strrchr(_profileName, '.');
		if(p) *p = 0;
	}

	//load XER file
	char * buf = malloc(CERT_MAX_SIZE);
	char * ebuf;
	EtsiTs103097Certificate_t * cert = NULL;

	ebuf = cstrnload(buf, CERT_MAX_SIZE, argv[1]);
	if(ebuf == NULL){
		fprintf(stderr, "%s: Certificate profile not found\n", argv[1]);
		return -1;
	}

	FILE * f;
	asn_dec_rval_t rc_d;
	asn_enc_rval_t rc_e;

	asn_TYPE_operation_t issuerOps = *asn_DEF_IssuerIdentifier.op;
	asn_DEF_IssuerIdentifier.op = &issuerOps;

	asn_TYPE_operation_t tbsOps = *asn_DEF_ToBeSignedCertificate.op;
	asn_DEF_ToBeSignedCertificate.op = &tbsOps;

	asn_TYPE_operation_t signatureOps = *asn_DEF_Signature.op;
	asn_DEF_Signature.op = &signatureOps;

	issuerOps.xer_decoder = IssuerIdentifier_xer_decoder;
	tbsOps.oer_encoder = ToBeSignedCertificate_oer_encoder;
	signatureOps.oer_encoder = Signature_oer_encoder;


	rc_d = asn_decode(NULL, ATS_BASIC_XER, &asn_DEF_EtsiTs103097Certificate, (void**)&cert, buf, ebuf - buf);
	if (rc_d.code != RC_OK){
		fprintf(stderr, "%s: failed to load at position %d\n  %.30s\n", _profileName, rc_d.consumed, buf + rc_d.consumed);
		return -1;
	}

	//check signer
	if (!_signerName && cert->issuer.present != IssuerIdentifier_PR_self){
		fprintf(stderr, "%s: unknown signer\n", argv[1]);
		return -1;
	}
	if (_signerName){
		PublicVerificationKey_PR hashType = PublicVerificationKey_PR_NOTHING;

		cvstrncpy(buf, CERT_MAX_SIZE, _searchPath, "/", _signerName, ".oer", NULL);
		ebuf = cstrnload(buf, CERT_MAX_SIZE, buf);
		if (ebuf == NULL){
			fprintf(stderr, "%s: signer certificate not found", _signerName);
			return -1;
		}
		// decode it to detect hash algorythm
		asn_dec_rval_t rc_d;
		EtsiTs103097Certificate_t * signer = NULL;
		rc_d = asn_decode(NULL, ATS_BASIC_OER, &asn_DEF_EtsiTs103097Certificate, (void**)&signer, buf, ebuf - buf);
		if (rc_d.code != RC_OK){
			fprintf(stderr, "%s: failed to load signer certificate at position %d\n  %.30s\n", _signerName, rc_d.consumed, buf + rc_d.consumed);
			return -1;
		}
		switch (signer->toBeSigned.verifyKeyIndicator.present){
		case VerificationKeyIndicator_PR_verificationKey:
			hashType = signer->toBeSigned.verifyKeyIndicator.choice.verificationKey.present;
			break;
		case VerificationKeyIndicator_PR_reconstructionValue:
			hashType = PublicVerificationKey_PR_ecdsaNistP256;
			break;
		}

		switch (hashType){
		case PublicVerificationKey_PR_ecdsaBrainpoolP256r1:
		case PublicVerificationKey_PR_ecdsaNistP256:
			if (cert->issuer.present == IssuerIdentifier_PR_NOTHING)
				cert->issuer.present = IssuerIdentifier_PR_sha256AndDigest;
			sha256_calculate(_signerHash, buf, ebuf - buf);
			_signerHashLength = sha256_hash_size;
			OCTET_STRING_fromBuf(&cert->issuer.choice.sha256AndDigest, &_signerHash[sha256_hash_size-8], 8);
			break;
		case PublicVerificationKey_PR_ecdsaBrainpoolP384r1:
			if (cert->issuer.present == IssuerIdentifier_PR_NOTHING)
				cert->issuer.present = IssuerIdentifier_PR_sha384AndDigest;
			sha384_calculate(_signerHash, buf, ebuf - buf);
			_signerHashLength = sha384_hash_size;
			OCTET_STRING_fromBuf(&cert->issuer.choice.sha384AndDigest, &_signerHash[sha384_hash_size - 8], 8);
			break;
		}
		if (signer){
			ASN_STRUCT_FREE(asn_DEF_EtsiTs103097Certificate, signer);
		}
	}

	// generate keys if necessary
	// buf = name of private key file
	cvstrncpy(buf, CERT_MAX_SIZE, _keyPath, "/", _profileName, EXT_VKEY, NULL);
	if (_force){
		remove(buf);
	}
	switch (cert->toBeSigned.verifyKeyIndicator.present){
	case VerificationKeyIndicator_PR_verificationKey:
		switch (cert->toBeSigned.verifyKeyIndicator.choice.verificationKey.present){
		case PublicVerificationKey_PR_ecdsaNistP256:
			fill_curve_point_eccP256(&cert->toBeSigned.verifyKeyIndicator.choice.verificationKey.choice.ecdsaNistP256, ecies_nistp256, buf);
			break;
		case PublicVerificationKey_PR_ecdsaBrainpoolP256r1:
			fill_curve_point_eccP256(&cert->toBeSigned.verifyKeyIndicator.choice.verificationKey.choice.ecdsaBrainpoolP256r1, ecies_brainpoolp256r, buf);
			break;
		case PublicVerificationKey_PR_ecdsaBrainpoolP384r1:
			fill_curve_point_eccP384(&cert->toBeSigned.verifyKeyIndicator.choice.verificationKey.choice.ecdsaBrainpoolP384r1, ecies_brainpoolp384r, buf);
			break;
		default:
			fprintf(stderr, "Unknown verification key curve type\n");
			return -1;
		}
		break;
	case VerificationKeyIndicator_PR_reconstructionValue:
		fprintf(stderr, "TODO: reconstruction value generation is unsupported\n");
	case VerificationKeyIndicator_PR_NOTHING:
	default:
		break;
	}
	if (cert->toBeSigned.encryptionKey){
		cvstrncpy(buf, CERT_MAX_SIZE, _keyPath, "/", _profileName, EXT_EKEY, NULL);
		if (_force){
			remove(buf);
		}
		switch (cert->toBeSigned.encryptionKey->publicKey.present){
		case BasePublicEncryptionKey_PR_NOTHING:
			cert->toBeSigned.encryptionKey->publicKey.present = BasePublicEncryptionKey_PR_eciesNistP256;
		case BasePublicEncryptionKey_PR_eciesNistP256:
			fill_curve_point_eccP256(&cert->toBeSigned.encryptionKey->publicKey.choice.eciesNistP256, ecies_nistp256, buf);
			break;
		case BasePublicEncryptionKey_PR_eciesBrainpoolP256r1:
			fill_curve_point_eccP256(&cert->toBeSigned.encryptionKey->publicKey.choice.eciesBrainpoolP256r1, ecies_brainpoolp256r, buf);
			break;
		default:
			break;
		}
	}

	cvstrncpy(buf, CERT_MAX_SIZE, _outPath, "/", _profileName, EXT_CERT, NULL);
	f = fopen(buf, "wb");
	if (f == NULL){
		perror(buf);
		return -1;
	}

	_cert = cert;
	// Encode as OER
	rc_e = asn_encode_to_buffer(NULL, ATS_CANONICAL_OER, &asn_DEF_EtsiTs103097Certificate, cert, buf, CERT_MAX_SIZE);
	if (rc_e.encoded <0){
		fprintf(stderr, "%s: OER encoding failed for %s\n", _profileName, rc_e.failed_type->name);
		return -1;
	}

	fwrite(buf, 1, rc_e.encoded, f);
	fclose(f);

	if(cert){
		ASN_STRUCT_FREE(asn_DEF_EtsiTs103097Certificate, cert);
	}
	return 0;
}

static int is_P256CurvePoint_empty(EccP256CurvePoint_t* point)
{
	switch (point->present){
	case EccP256CurvePoint_PR_x_only:
	case EccP256CurvePoint_PR_compressed_y_0:
	case EccP256CurvePoint_PR_compressed_y_1:
		return point->choice.x_only.size < 32;
	case EccP256CurvePoint_PR_uncompressedP256:
		return point->choice.uncompressedP256.x.size < 32 || point->choice.uncompressedP256.y.size < 32;
	default:
		break;
	}
	return 1;
}
static void fill_curve_point_eccP256(EccP256CurvePoint_t* point, ecc_curve_id curveType, char * keyPath)
{
	void * key;
	char x[32], y[32];
	int compressed_y;
	key = ecc_key_private_load(keyPath, curveType);
	if (key == NULL){
		key = ecc_key_gen(curveType);
		ecc_key_private_save(key, keyPath, _outKeyFormat);
		strcat(keyPath, EXT_PUB);
		ecc_key_public_save(key, keyPath, _outKeyFormat);
	}
	ecc_key_public(key, x, y, &compressed_y);
	OCTET_STRING_fromBuf(&point->choice.x_only, x, 32);
	if (point->present == EccP256CurvePoint_PR_uncompressedP256){
		OCTET_STRING_fromBuf(&point->choice.uncompressedP256.y, y, 32);
	}
	else if (point->present == EccP256CurvePoint_PR_compressed_y_0 || point->present == EccP256CurvePoint_PR_compressed_y_1){
		point->present = compressed_y ? EccP256CurvePoint_PR_compressed_y_1 : EccP256CurvePoint_PR_compressed_y_0;
	}
	ecc_key_free(key);
}

static int is_P384CurvePoint_empty(EccP384CurvePoint_t* point){
	switch (point->present){
	case EccP256CurvePoint_PR_x_only:
	case EccP256CurvePoint_PR_compressed_y_0:
	case EccP256CurvePoint_PR_compressed_y_1:
		return point->choice.x_only.size < 48;
	case EccP256CurvePoint_PR_uncompressedP256:
		return point->choice.uncompressedP384.x.size < 48 || point->choice.uncompressedP384.y.size < 48;
	default:
		break;
	}
	return 1;
}

static void fill_curve_point_eccP384(EccP384CurvePoint_t* point, ecc_curve_id curveType, char * keyPath)
{
	void * key;
	char x[48], y[48];
	int compressed_y;
	key = ecc_key_private_load(keyPath, curveType);
	if (key == NULL){
		key = ecc_key_gen(curveType);
		ecc_key_private_save(key, keyPath, _outKeyFormat);
		strcat(keyPath, EXT_PUB);
		ecc_key_public_save(key, keyPath, _outKeyFormat);
	}
	ecc_key_public(key, x, y, &compressed_y);
	OCTET_STRING_fromBuf(&point->choice.x_only, x, 48);
	if (point->present == EccP384CurvePoint_PR_uncompressedP384){
		OCTET_STRING_fromBuf(&point->choice.uncompressedP384.y, y, 48);
	}
	else if (point->present == EccP256CurvePoint_PR_compressed_y_0 || point->present == EccP256CurvePoint_PR_compressed_y_1){
		point->present = compressed_y ? EccP256CurvePoint_PR_compressed_y_1 : EccP256CurvePoint_PR_compressed_y_0;
	}
	ecc_key_free(key);
}