#!/usr/bin/env python3
"""Resolve a TiDB Cloud Serverless shared-pool name into parent routing identities.

Accepts either a bare pool name (``us-west-2-f01``) or a full Grafana
``control_plane_info`` value (``us-west-2-f01/aws-us-west-2/devtier-infra``)
and prints every derived routing identity used by the serverless parent
metrics and logs APIs.
"""

from __future__ import annotations

import argparse
import json
import re
import shlex
import sys
from dataclasses import dataclass

VENDOR_CODES = {"f": "aws", "a": "alicloud"}
DEFAULT_APPLICATION = "devtier-infra"
ROUTING_SEGMENT_RE = re.compile(r"[A-Za-z0-9._:+-]+\Z")


@dataclass(frozen=True)
class PoolRouting:
    pool: str
    vendor: str
    vendor_region: str
    application: str

    @property
    def control_plane_info(self) -> str:
        return f"{self.pool}/{self.vendor_region}/{self.application}"

    @property
    def promql_control_plane_info(self) -> str:
        return f"{self.pool}/{self.vendor_region}"

    @property
    def loki_datasource(self) -> str:
        return f"Loki-{self.vendor_region}"

    @property
    def prod_loki_cluster(self) -> str:
        return f"prod-{self.pool}"

    def as_dict(self) -> dict[str, str]:
        return {
            "pool": self.pool,
            "vendor": self.vendor,
            "vendor_region": self.vendor_region,
            "application": self.application,
            "control_plane_info": self.control_plane_info,
            "promql_control_plane_info": self.promql_control_plane_info,
            "loki_datasource": self.loki_datasource,
            "prod_loki_cluster": self.prod_loki_cluster,
        }


def vendor_and_region_from_pool(pool: str) -> tuple[str, str]:
    segments = pool.split("-")
    if (
        len(segments) != 4
        or not all(segments)
        or not all(ROUTING_SEGMENT_RE.fullmatch(segment) for segment in segments)
    ):
        raise ValueError(
            f"pool name must have four dash-separated segments (xx-xx-xx-xx): {pool!r}"
        )
    code = segments[3][0]
    vendor = VENDOR_CODES.get(code)
    if vendor is None:
        known = ", ".join(f"{key}={value}" for key, value in sorted(VENDOR_CODES.items()))
        raise ValueError(f"unknown vendor code {code!r} in pool {pool!r} (known: {known})")
    return vendor, f"{vendor}-{'-'.join(segments[:3])}"


def resolve(pool_or_control_plane_info: str, application: str | None = None) -> PoolRouting:
    parts = pool_or_control_plane_info.split("/")
    if len(parts) == 1 and parts[0]:
        pool = parts[0]
        vendor, vendor_region = vendor_and_region_from_pool(pool)
        resolved_application = DEFAULT_APPLICATION if application is None else application
    elif len(parts) == 3 and all(parts):
        pool, given_vendor_region, given_application = parts
        vendor, vendor_region = vendor_and_region_from_pool(pool)
        if given_vendor_region != vendor_region:
            raise ValueError(
                f"vendor-region {given_vendor_region!r} does not match pool {pool!r} "
                f"(expected {vendor_region!r})"
            )
        resolved_application = given_application if application is None else application
    else:
        raise ValueError(
            "value must be a bare pool name (xx-xx-xx-xx) or the full Grafana "
            "routing value <pool>/<vendor-region>/<application>: "
            f"{pool_or_control_plane_info!r}"
        )
    if not ROUTING_SEGMENT_RE.fullmatch(resolved_application):
        raise ValueError(
            "application must contain only letters, digits, '.', '_', '-', ':', or '+'"
        )
    return PoolRouting(pool, vendor, vendor_region, resolved_application)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Resolve a Serverless pool name into parent metrics/logs routing identities."
    )
    parser.add_argument(
        "value",
        help="bare pool name (us-west-2-f01) or full control_plane_info "
        "(us-west-2-f01/aws-us-west-2/devtier-infra)",
    )
    parser.add_argument(
        "--application",
        default=None,
        help=f"override the application ID (default: {DEFAULT_APPLICATION})",
    )
    parser.add_argument("--json", action="store_true", help="print one JSON object")
    args = parser.parse_args(argv)

    try:
        routing = resolve(args.value, args.application)
    except ValueError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    if args.json:
        print(json.dumps(routing.as_dict(), indent=2))
    else:
        for key, value in routing.as_dict().items():
            print(f"{key}={shlex.quote(value)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())