"""Direct-to-S3 upload with AWS SigV4, using only the standard library.

The studio hands out temporary AWS credentials via
``POST /v2/studio/storage/temporary_token``; the browser then PUTs the raw bytes
straight to S3 (``https://{bucket}.{host}/{key}``).  No boto3 required.
"""

import datetime
import hashlib
import hmac
import urllib.request

REGION = "us-west-2"
SERVICE = "s3"


def _sign(key, message):
    return hmac.new(key, message.encode("utf-8"), hashlib.sha256).digest()


def upload_bytes(sts, data, content_type="application/octet-stream", timeout=120):
    """PUT ``data`` to the S3 object described by the STS response.

    ``sts`` is the ``data`` block of ``/v2/studio/storage/temporary_token``:
    ``sts_ak``, ``sts_sk``, ``session_token``, ``host``, ``resource_bucket``,
    ``resource_uri``.
    """
    access_key = sts["sts_ak"]
    secret_key = sts["sts_sk"]
    session_token = sts["session_token"]
    host = sts.get("host", "s3.us-west-2.amazonaws.com")
    bucket = sts["resource_bucket"]
    key = sts["resource_uri"]

    now = datetime.datetime.now(datetime.timezone.utc)
    amz_date = now.strftime("%Y%m%dT%H%M%SZ")
    date_stamp = now.strftime("%Y%m%d")

    host_header = f"{bucket}.{host}"
    payload_hash = hashlib.sha256(data).hexdigest()

    canonical_headers = (
        f"host:{host_header}\n"
        f"x-amz-content-sha256:{payload_hash}\n"
        f"x-amz-date:{amz_date}\n"
        f"x-amz-security-token:{session_token}\n"
    )
    signed_headers = "host;x-amz-content-sha256;x-amz-date;x-amz-security-token"
    canonical_request = f"PUT\n/{key}\n\n{canonical_headers}\n{signed_headers}\n{payload_hash}"

    scope = f"{date_stamp}/{REGION}/{SERVICE}/aws4_request"
    string_to_sign = (
        f"AWS4-HMAC-SHA256\n{amz_date}\n{scope}\n"
        f"{hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()}"
    )

    k_date = _sign(("AWS4" + secret_key).encode("utf-8"), date_stamp)
    k_region = _sign(k_date, REGION)
    k_service = _sign(k_region, SERVICE)
    k_signing = _sign(k_service, "aws4_request")
    signature = hmac.new(k_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()

    authorization = (
        f"AWS4-HMAC-SHA256 Credential={access_key}/{scope}, "
        f"SignedHeaders={signed_headers}, Signature={signature}"
    )

    request = urllib.request.Request(
        f"https://{host_header}/{key}",
        data=data,
        method="PUT",
        headers={
            "Host": host_header,
            "x-amz-date": amz_date,
            "x-amz-content-sha256": payload_hash,
            "x-amz-security-token": session_token,
            "Authorization": authorization,
            "Content-Type": content_type,
        },
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return response.status in (200, 204)