"""
Prosody mod_http_upload_external -> S3 handler.

Implements the same HMAC protocol as prosody-filer but stores files in S3.

Flow:
  1. XMPP client asks Prosody for an upload slot
  2. Prosody generates a signed PUT URL pointing here
  3. Client PUTs the file; this handler verifies the HMAC and uploads to S3
  4. On GET, handler generates a presigned S3 URL and redirects
"""

import hashlib
import hmac as hmac_mod
import logging
import os

import boto3
from flask import Flask, Response, redirect, request

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("s3-upload-handler")

SECRET = os.environ["UPLOAD_SECRET"]
S3_BUCKET = os.environ["S3_BUCKET"]
S3_REGION = os.environ.get("S3_REGION", "us-east-1")
S3_ENDPOINT = os.environ.get("S3_ENDPOINT")
PRESIGN_EXPIRE = int(os.environ.get("PRESIGN_EXPIRE", "3600"))

s3_kwargs = {
    "region_name": S3_REGION,
}
if S3_ENDPOINT:
    s3_kwargs["endpoint_url"] = S3_ENDPOINT

s3 = boto3.client("s3", **s3_kwargs)


def verify_token(path: str, size: str, content_type: str, token: str) -> bool:
    """Verify HMAC-SHA256 token from Prosody mod_http_upload_external."""
    message = f"{path} {size} {content_type}"
    expected = hmac_mod.new(
        SECRET.encode(), message.encode(), hashlib.sha256
    ).hexdigest()
    return hmac_mod.compare_digest(expected, token)


@app.route("/upload/<path:file_path>", methods=["PUT"])
def upload(file_path):
    token = request.args.get("v", "")
    content_type = request.args.get("t", "application/octet-stream")
    file_size = request.args.get("s", "0")

    hmac_path = f"upload/{file_path}"

    if not verify_token(hmac_path, file_size, content_type, token):
        log.warning("HMAC verification failed for %s", hmac_path)
        return Response("Forbidden", status=403)

    body = request.get_data()

    if len(body) != int(file_size):
        return Response("Content length mismatch", status=400)

    s3.put_object(
        Bucket=S3_BUCKET,
        Key=file_path,
        Body=body,
        ContentType=content_type,
    )

    log.info("Uploaded %s (%s bytes) to s3://%s/%s", file_path, file_size, S3_BUCKET, file_path)
    return Response(status=201)


@app.route("/upload/<path:file_path>", methods=["GET", "HEAD"])
def download(file_path):
    try:
        url = s3.generate_presigned_url(
            "get_object",
            Params={"Bucket": S3_BUCKET, "Key": file_path},
            ExpiresIn=PRESIGN_EXPIRE,
        )
    except Exception:
        log.exception("Failed to generate presigned URL for %s", file_path)
        return Response("Not found", status=404)

    return redirect(url, code=302)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5050)