# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "requests",
#     "boto3",
# ]
# ///
import argparse
import json
import sys

import boto3

from lib import (
    O11yClient,
    get_auth0_token,
    load_auth0_secret_string_azure,
    load_auth0_secret_string_aws,
    load_auth0_secret_string_alicloud,
    parse_account_id_from_role_arn,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create an o11y cluster and wait for it to become ACTIVE"
    )
    parser.add_argument("--cloud", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--logical-region-name", required=True)
    parser.add_argument("--o11y-endpoint", required=True)
    parser.add_argument("--tenant-id", required=True)
    parser.add_argument("--project-id", required=True)
    parser.add_argument("--biz-type", required=True)
    parser.add_argument("--rule-template-id", required=True)
    parser.add_argument("--vector-role-arn", default=None)
    parser.add_argument("--agent-access-account", default=None)
    parser.add_argument("--key-vault-name", default=None)
    parser.add_argument("--dataplane-manager-object-id", default=None)
    parser.add_argument("--controlplane-subscription-id", default=None)
    parser.add_argument("--auth0-secret-id", required=True)
    parser.add_argument("--resource-id", required=True)
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--interval", type=int, default=10)
    return parser.parse_args()


def store_cluster_id(bucket: str, key: str, cluster_data: dict) -> None:
    boto3.client("s3").put_object(
        Bucket=bucket,
        Key=key,
        Body=json.dumps(cluster_data).encode("utf-8"),
        ContentType="application/json",
    )


def find_cluster_by_seed_tag(clusters: list[dict], seed_tag: str) -> dict:
    for cluster in clusters:
        if seed_tag in cluster.get("tags", []):
            return cluster
    raise RuntimeError(
        f"No o11y cluster found with seed tag '{seed_tag}' after creation"
    )


if __name__ == "__main__":
    args = parse_args()
    biz_type = args.biz_type
    seed_tag = (
        args.resource_id
    )  # Use resource_id as the seed tag to identify the cluster

    print(
        "===================== start o11y cluster creation =====================",
        file=sys.stderr,
    )

    if args.cloud == "alicloud":
        token = get_auth0_token(
            load_auth0_secret_string_alicloud(args.auth0_secret_id, args.agent_access_account, args.region)
        )
    elif args.cloud == "azure":
        token = get_auth0_token(
            load_auth0_secret_string_azure(
                args.auth0_secret_id,
                args.key_vault_name,
                args.controlplane_subscription_id,
            )
        )
    else:
        profile_name = parse_account_id_from_role_arn(args.vector_role_arn)
        token = get_auth0_token(
            load_auth0_secret_string_aws(args.auth0_secret_id, profile_name, args.region)
        )
    client = O11yClient(args.o11y_endpoint, token, args.tenant_id, biz_type)

    print("Creating o11y cluster...", file=sys.stderr)
    create_body = {
        "vendor": args.cloud,
        "region": args.region,
        "tenantID": args.tenant_id,
        "projectID": args.project_id,
        "applicationID": args.logical_region_name,
        "bizType": biz_type,
        "version": "v1",
        "shareMode": "shared",
        "dataPropertyRights": 1,
        "ruleTemplateID": args.rule_template_id,
        "tags": [seed_tag],
    }
    if args.cloud == "alicloud":
        create_body["agentAccessAccount"] = args.agent_access_account
    elif args.cloud == "azure":
        create_body["agentAccessARN"] = args.dataplane_manager_object_id
        create_body["agentManagerARN"] = args.controlplane_subscription_id
    else:
        create_body["agentAccessARN"] = args.vector_role_arn
    create_resp = client.create_cluster(create_body)
    print(f"create_cluster_resp: {create_resp}", file=sys.stderr)
    if create_resp.get("code") != 0:
        raise RuntimeError(
            f"Failed to create o11y cluster, code: {create_resp.get('code')}"
        )

    clusters = client.list_clusters(
        {
            "bizType": biz_type,
            "region": args.region,
            "tenantID": args.tenant_id,
            "vendor": args.cloud,
            "projectID": args.project_id,
            "applicationID": args.logical_region_name,
        }
    )
    if not clusters:
        raise RuntimeError("No o11y clusters found after creation")

    cluster = find_cluster_by_seed_tag(clusters, seed_tag)
    cluster_id = cluster["id"]
    print(f"o11y_cluster_id: {cluster_id}", file=sys.stderr)

    # Wait for the cluster to become ACTIVE
    cluster_data = client.wait_for_active(
        cluster_id,
        args.cloud,
        args.region,
        timeout=args.timeout,
        interval=args.interval,
    )
    store_cluster_id("autonomous-deployment-kv", args.resource_id, cluster_data)
    print(
        f"stored o11y cluster data in s3://autonomous-deployment-kv/{args.resource_id}",
        file=sys.stderr,
    )
    print(f"o11y cluster ready: {cluster_id}", file=sys.stderr)
    print(
        "===================== o11y cluster creation done =====================",
        file=sys.stderr,
    )