// Copyright (c) 2019-2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module s3

import crypto.hmac
import crypto.sha256
import time

// service_name is the SigV4 service identifier baked into the signing key.
pub const service_name = 's3'

// algo is the SigV4 algorithm marker S3 expects.
pub const algo = 'AWS4-HMAC-SHA256'

// unsigned_payload is the magic string used in `x-amz-content-sha256` when
// the payload is not pre-hashed. Used by the single-shot put path to avoid
// buffering / re-scanning the entire body just to sign it.
pub const unsigned_payload = 'UNSIGNED-PAYLOAD'

// empty_sha256 is `sha256("")` precomputed. Used for HEAD/GET/DELETE where
// there is no body and we want a real hash for stricter S3 endpoints.
pub const empty_sha256 = 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'

// SignRequest describes the request to sign. The signer is intentionally
// agnostic of HTTP option semantics (ACL, storage class, …) — the caller
// pre-fills `extra_headers` with whatever it intends to send. That keeps the
// signer pure and easy to test against published SigV4 reference vectors.
pub struct SignRequest {
pub:
	method        string @[required] // GET, PUT, POST, DELETE, HEAD
	path          string @[required] // canonical URI (already URI-encoded, kept as-is)
	query         string            // canonical query string (sorted, encoded; no leading '?')
	payload_hash  string            // hex SHA-256 of body or `unsigned_payload`
	extra_headers map[string]string // lowercase keys, raw values — will be added to canonical/signed set
	sign_time     time.Time         // when omitted (Time{}) we use time.utc()
}

// SignedRequest is the output of header signing.
pub struct SignedRequest {
pub:
	method        string            // canonical HTTP method that was signed (GET, PUT, …)
	url           string            // scheme://host<extra_path>/<bucket>/<key>?<query>
	host          string            // host[:port], suitable for the Host header
	amz_date      string            // YYYYMMDDTHHMMSSZ
	authorization string            // ready-to-send Authorization header value
	headers       map[string]string // ALL headers that MUST be sent for the signature to verify
}

// sign_request builds the SigV4 Authorization header for an HTTP request.
//
// The returned `SignedRequest.headers` is the full set the caller must send
// (excluding `Content-Length`, which the HTTP client adds itself). Adding,
// removing, or mutating a header after signing will break the signature.
pub fn sign_request(creds Credentials, req SignRequest) !SignedRequest {
	creds.validate()!
	method := normalize_method(req.method) or {
		return new_error('InvalidMethod',
			'Method must be GET, PUT, POST, DELETE or HEAD; got: ${req.method}')
	}
	host := canonical_host(creds, '')
	if host == '' {
		return new_error('InvalidEndpoint',
			'No endpoint and no bucket provided — cannot determine host')
	}
	region := creds.resolved_region()
	now := if req.sign_time.year == 0 { time.utc() } else { req.sign_time }
	amz_date := format_amz_date(now)
	amz_day := amz_date[..8]
	payload := if req.payload_hash == '' { unsigned_payload } else { req.payload_hash }

	// Build the headers set. We always sign host + amz date + content sha
	// because S3 requires them. Reject CRLF before they reach the wire.
	mut headers := map[string]string{}
	for k, v in req.extra_headers {
		if contains_crlf(v) || contains_crlf(k) {
			return new_error('InvalidHeader',
				'Header "${k}" contains CR/LF — refused (header injection guard)')
		}
		headers[k.to_lower()] = v
	}
	headers['host'] = host
	headers['x-amz-content-sha256'] = payload
	headers['x-amz-date'] = amz_date
	if creds.session_token != '' {
		headers['x-amz-security-token'] = creds.session_token
	}

	signed_header_names := sorted_keys(headers).join(';')
	canonical := build_canonical_request(method, req.path, req.query, headers, signed_header_names,
		payload)
	credential_scope := '${amz_day}/${region}/${service_name}/aws4_request'
	string_to_sign := '${algo}\n${amz_date}\n${credential_scope}\n${sha256_hex(canonical.bytes())}'
	signature := to_hex_lower(hmac_sha256(derive_signing_key(creds.secret_access_key, amz_day,
		region, service_name), string_to_sign.bytes()))
	authorization := '${algo} Credential=${creds.access_key_id}/${credential_scope}, SignedHeaders=${signed_header_names}, Signature=${signature}'
	headers['authorization'] = authorization

	scheme := creds.scheme()
	mut url := '${scheme}://${host}${req.path}'
	if req.query != '' {
		url += '?' + req.query
	}
	return SignedRequest{
		method:        method
		url:           url
		host:          host
		amz_date:      amz_date
		authorization: authorization
		headers:       headers
	}
}

// PresignRequest describes a presigned URL to generate. The output is a
// self-contained URL — no extra headers are required at request time.
pub struct PresignRequest {
pub:
	method      string @[required]
	path        string @[required] // canonical URI (already URI-encoded)
	expires_in  int = 86400 // 1..604800 seconds (SigV4 hard limit is 7 days)
	extra_query map[string]string // additional signed query params (e.g. response-content-type, x-amz-acl)
	sign_time   time.Time
}

// presign_url returns a fully-formed `https://...` URL signed via SigV4
// query-string parameters. The URL is valid until `now + expires_in`.
pub fn presign_url(creds Credentials, req PresignRequest) !string {
	creds.validate()!
	method := normalize_method(req.method) or {
		return new_error('InvalidMethod',
			'Method must be GET, PUT, POST, DELETE or HEAD; got: ${req.method}')
	}
	if req.expires_in < 1 || req.expires_in > 604800 {
		return new_error('InvalidExpiry', 'expires_in must be between 1 and 604800 seconds')
	}
	host := canonical_host(creds, '')
	if host == '' {
		return new_error('InvalidEndpoint',
			'No endpoint and no bucket provided — cannot determine host')
	}
	region := creds.resolved_region()
	now := if req.sign_time.year == 0 { time.utc() } else { req.sign_time }
	amz_date := format_amz_date(now)
	amz_day := amz_date[..8]
	credential := '${creds.access_key_id}/${amz_day}/${region}/${service_name}/aws4_request'

	// Required signed query parameters; we then merge in extras from the caller.
	mut params := map[string]string{}
	params['X-Amz-Algorithm'] = algo
	params['X-Amz-Credential'] = credential
	params['X-Amz-Date'] = amz_date
	params['X-Amz-Expires'] = req.expires_in.str()
	params['X-Amz-SignedHeaders'] = 'host'
	if creds.session_token != '' {
		params['X-Amz-Security-Token'] = creds.session_token
	}
	for k, v in req.extra_query {
		if contains_crlf(k) || contains_crlf(v) {
			return new_error('InvalidQueryParam', 'Query param "${k}" contains CR/LF — refused')
		}
		params[k] = v
	}

	canonical_query := canonical_query_string(params)
	headers := {
		'host': host
	}
	canonical := build_canonical_request(method, req.path, canonical_query, headers, 'host',
		unsigned_payload)
	credential_scope := '${amz_day}/${region}/${service_name}/aws4_request'
	string_to_sign := '${algo}\n${amz_date}\n${credential_scope}\n${sha256_hex(canonical.bytes())}'
	signature := to_hex_lower(hmac_sha256(derive_signing_key(creds.secret_access_key, amz_day,
		region, service_name), string_to_sign.bytes()))
	scheme := creds.scheme()
	return '${scheme}://${host}${req.path}?${canonical_query}&X-Amz-Signature=${signature}'
}

// build_canonical_request assembles the canonical request string per
// SigV4 §3.2. `path` and `query` MUST already be URI-encoded.
pub fn build_canonical_request(method string, path string, query string, headers map[string]string,
	signed_headers string, payload_hash string) string {
	sorted := sorted_keys(headers)
	mut lines := []string{cap: sorted.len}
	for k in sorted {
		lines << '${k}:${normalize_header_value(headers[k])}'
	}
	return '${method}\n${path}\n${query}\n${lines.join('\n')}\n\n${signed_headers}\n${payload_hash}'
}

// canonical_query_string sorts query params by key (then by value if the same
// key appears twice — we don't here) and joins them as `k=v&k=v` with values
// already URI-encoded.
pub fn canonical_query_string(params map[string]string) string {
	keys := sorted_keys(params)
	mut parts := []string{cap: keys.len}
	for k in keys {
		parts << '${uri_encode_query(k)}=${uri_encode_query(params[k])}'
	}
	return parts.join('&')
}

// derive_signing_key implements the four-step HMAC chain that produces the
// SigV4 signing key. The result is cacheable per (secret, date, region,
// service) tuple — left to the caller to memoize if needed.
pub fn derive_signing_key(secret string, date string, region string, service string) []u8 {
	k_date := hmac_sha256(('AWS4' + secret).bytes(), date.bytes())
	k_region := hmac_sha256(k_date, region.bytes())
	k_service := hmac_sha256(k_region, service.bytes())
	return hmac_sha256(k_service, 'aws4_request'.bytes())
}

// format_amz_date returns the basic ISO-8601 timestamp used by SigV4
// (`YYYYMMDDTHHMMSSZ`). `t` is assumed to be in UTC.
pub fn format_amz_date(t time.Time) string {
	return '${t.year:04d}${t.month:02d}${t.day:02d}T${t.hour:02d}${t.minute:02d}${t.second:02d}Z'
}

// canonical_host resolves the on-the-wire host. Path-style addressing is
// the default; virtual-hosted style uses `<bucket>.<endpoint-host>`. If no
// endpoint is configured we fall back to `s3.<region>.amazonaws.com` (the
// canonical default for clients pointed at AWS S3 itself).
pub fn canonical_host(creds Credentials, bucket_override string) string {
	bucket := if bucket_override != '' { bucket_override } else { creds.bucket }
	host := creds.host_only()
	if host == '' {
		region := creds.resolved_region()
		if creds.virtual_hosted_style {
			if bucket == '' {
				return ''
			}
			return '${bucket}.s3.${region}.amazonaws.com'
		}
		return 's3.${region}.amazonaws.com'
	}
	if creds.virtual_hosted_style && bucket != '' {
		return '${bucket}.${host}'
	}
	return host
}

// normalize_method uppercases and validates the HTTP method.
pub fn normalize_method(m string) ?string {
	up := m.to_upper()
	return match up {
		'GET', 'PUT', 'POST', 'DELETE', 'HEAD' { up }
		else { none }
	}
}

// normalize_header_value collapses runs of internal whitespace to a single
// space and trims leading/trailing whitespace, per SigV4 §3.2.2 step 3.
fn normalize_header_value(v string) string {
	if v == '' {
		return v
	}
	mut out := []u8{cap: v.len}
	mut prev_space := false
	mut started := false
	for b in v.bytes() {
		if b == ` ` || b == `\t` {
			if started {
				prev_space = true
			}
			continue
		}
		if prev_space {
			out << ` `
			prev_space = false
		}
		out << b
		started = true
	}
	return out.bytestr()
}

// sha256_hex returns the lowercase hex digest of `data`.
@[inline]
pub fn sha256_hex(data []u8) string {
	return to_hex_lower(sha256.sum256(data))
}

// hmac_sha256 wraps `crypto.hmac.new` for HMAC-SHA-256 with the V stdlib's
// type signature (a hash function that returns `[]u8`).
@[inline]
pub fn hmac_sha256(key []u8, data []u8) []u8 {
	return hmac.new(key, data, sha256.sum256, sha256.block_size)
}

// sorted_keys returns the keys of `m` in lexical order.
fn sorted_keys(m map[string]string) []string {
	mut keys := m.keys()
	keys.sort()
	return keys
}