#!/usr/bin/env python3
"""Seed TiDB-importable CSV and SQL fixtures into the Date Integration OSS import prefixes.

For premium import E2E, the private test OSS bucket must contain data under both
import/csv/ and import/sql/. This runs on the ECS instance and reuses the ECS
RAM-role credentials (the same path the OSS smoke uses), so no static AK/SK is
committed. The payload manifest is base64-encoded and passed via --manifest-b64
(the shell reads the tracked scripts/seeddata/import/{csv,sql}/ files).
"""

import argparse
import base64
import hashlib
import hmac
import json
import urllib.parse
import urllib.request
from email.utils import formatdate


def load_credentials(role_name: str) -> dict[str, str]:
    role = urllib.parse.quote(role_name, safe="")
    url = f"http://100.100.100.200/latest/meta-data/ram/security-credentials/{role}"
    with urllib.request.urlopen(url, timeout=10) as response:
        credentials = json.load(response)
    if credentials.get("Code") != "Success":
        raise RuntimeError("ECS RAM-role metadata did not return Success")
    return credentials


def object_request(
    method: str,
    bucket: str,
    endpoint: str,
    object_key: str,
    credentials: dict[str, str],
    body: bytes | None = None,
) -> bytes:
    date = formatdate(usegmt=True)
    content_type = "application/octet-stream" if body is not None else ""
    token = credentials["SecurityToken"]
    canonical_headers = f"x-oss-security-token:{token}\n"
    canonical_resource = f"/{bucket}/{object_key}"
    string_to_sign = (
        f"{method}\n\n{content_type}\n{date}\n"
        f"{canonical_headers}{canonical_resource}"
    )
    digest = hmac.new(
        credentials["AccessKeySecret"].encode(),
        string_to_sign.encode(),
        hashlib.sha1,
    ).digest()
    signature = base64.b64encode(digest).decode()
    encoded_key = urllib.parse.quote(object_key, safe="/")
    request = urllib.request.Request(
        f"http://{bucket}.{endpoint}/{encoded_key}",
        data=body,
        method=method,
        headers={
            "Authorization": f"OSS {credentials['AccessKeyId']}:{signature}",
            "Date": date,
            "x-oss-security-token": token,
            **({"Content-Type": content_type} if content_type else {}),
        },
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        return response.read()


def decode_manifest(manifest_b64: str) -> list[tuple[str, bytes]]:
    try:
        manifest = json.loads(base64.b64decode(manifest_b64, validate=True))
    except Exception as exc:  # noqa: BLE001 - surface a concise decode error
        raise SystemExit(f"invalid --manifest-b64: {exc}") from exc

    if not isinstance(manifest, list) or not manifest:
        raise SystemExit("--manifest-b64 must contain a non-empty JSON list")

    decoded: list[tuple[str, bytes]] = []
    for entry in manifest:
        if not isinstance(entry, dict):
            raise SystemExit("each seed manifest entry must be a JSON object")
        object_key = entry.get("object_key")
        content_b64 = entry.get("content_b64")
        if (
            not isinstance(object_key, str)
            or not object_key.startswith("import/")
            or not isinstance(content_b64, str)
        ):
            raise SystemExit(
                "each seed entry must contain an import/ object_key and content_b64"
            )
        try:
            content = base64.b64decode(content_b64, validate=True)
        except Exception as exc:  # noqa: BLE001 - surface a concise decode error
            raise SystemExit(f"invalid content for {object_key}: {exc}") from exc
        decoded.append((object_key, content))
    return decoded


def seed_objects(
    bucket: str,
    endpoint: str,
    credentials: dict[str, str],
    entries: list[tuple[str, bytes]],
) -> None:
    if not entries:
        raise SystemExit("seed manifest is empty")
    for object_key, content in entries:
        object_request("PUT", bucket, endpoint, object_key, credentials, content)
        downloaded = object_request("GET", bucket, endpoint, object_key, credentials)
        if downloaded != content:
            raise RuntimeError(f"OSS import seed object did not round-trip: {object_key}")
        print(f"Seeded Alibaba Cloud OSS import data at oss://{bucket}/{object_key}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--bucket", required=True)
    parser.add_argument("--endpoint", required=True)
    parser.add_argument("--role-name", required=True)
    parser.add_argument(
        "--manifest-b64",
        required=True,
        help="base64-encoded JSON seed manifest from scripts/seeddata/import/",
    )
    args = parser.parse_args()

    entries = decode_manifest(args.manifest_b64)
    credentials = load_credentials(args.role_name)
    seed_objects(
        args.bucket,
        args.endpoint,
        credentials,
        entries,
    )


if __name__ == "__main__":
    main()