#!/usr/bin/env python3
"""Generate and enrich config/victoriametrics.json.

Reads the existing victoriametrics.json and preserves entries that already
have a valid datasource URL and labels.  Only regionIds not present in the
existing config go through the full generation pipeline (routing, Clusters
API enrichment, acs_build_info verification).  This allows manual overrides
to persist across regeneration and speeds up the process by skipping
already-resolved entries.

URL format:
    dataplaneAddress + "/internal/metrics/" + o11y_cluster_id

Pipeline
~~~~~~~~
For each regionId in config/regions.json:

1. **Route resolution** — resolve candidate dataplaneAddresses from
   meta.json routing (shared with the Loki script via routing.py).

2. **Clusters API enrichment** (optional, requires Auth0) — queries the
   Clusters API with ``bizType``, ``applicationID``, and optionally
   ``projectID`` filters to narrow results.  Matches clusters by
   ``internalReadUri`` prefix.  Strips any ``/api/v1/`` suffix from the
   returned URI.

   bizType mapping (derived from architecture + environment)::

     dedicated  + prod    → tidbcloud
     dedicated  + staging → staging-tidbcloud
     dedicated  + dev     → dev-tidbcloud
     nextgen    + prod    → prod-tidbcloud-nextgen
     nextgen    + staging → staging-tidbcloud-nextgen
     nextgen    + dev     → dev-tidbcloud-nextgen
     byoc       + *       → same as nextgen
     serverless + prod    → prod-devtier
     serverless + staging → staging-devtier
     serverless + dev     → dev-devtier

   applicationID (fixed infrastructure clusters)::

     dedicated  → seed
     serverless → devtier-infra
     nextgen    → (none — no fixed applicationID)

   projectID (serverless only)::

     projectID = regionName  (e.g. ``us-east-1-f01``)
     Returns exactly 1 cluster per shard, avoiding the 100-result cap.

3. **Verification + label enrichment** — queries ``acs_build_info`` on
   every entry to verify the URL works and extract the ``cluster`` label.
   On 4xx errors, falls back to the next candidate base URL (same logic
   as the Loki script).  Unresolved entries are printed as a summary at
   the end and cause a non-zero exit.

4. Final output = ``{regionId, datasource: {url}, labels: {container,
   cluster}}``.  Entries without a UUID or cluster label are excluded
   from the output.

Usage:
    make generate-vm
"""

import json
import os
import re
import sys
from pathlib import Path
from urllib.error import HTTPError
from urllib.parse import urlencode
from urllib.request import Request, urlopen

from routing import (
    parse_region_id,
    physical_region_keys,
    load_meta,
    resolve_candidate_bases,
    SUFFIX_RE,
)

UUID_RE = re.compile(
    r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.I,
)
CONTAINER = "agent-control-server"

HERE = Path(__file__).resolve().parent
PROJECT = HERE.parent

REGIONS_PATH = PROJECT / "config" / "regions.json"
META_PATH = (
    PROJECT / ".." / "deployments" / "globalserver"
    / "manifests" / "prod" / "meta.json"
)
VM_PATH = PROJECT / "config" / "victoriametrics.json"
VM_PATH_PREFIX = "/internal/metrics/"

# Auth0 configuration (from resolve_cluster_uri.py)
AUTH0_URL = os.getenv("AUTH0_URL", "https://tidb-soc2.us.auth0.com/oauth/token")
AUTH0_CLIENT_ID = os.getenv("AUTH0_CLIENT_ID")
AUTH0_CLIENT_SECRET = os.getenv("AUTH0_CLIENT_SECRET")
AUTH0_AUDIENCE = os.getenv("AUTH0_AUDIENCE", "https://tidb-soc2.us.auth0.com/api/v2/")

CLUSTERS_API_BASE = (
    "http://www.gs.us-west-2.aws.observability.tidbcloud.com/api/v1/clusters/"
)

BIZ_TYPE_MAP = {
    ("dedicated", "prod"): "tidbcloud",
    ("dedicated", "staging"): "staging-tidbcloud",
    ("dedicated", "dev"): "dev-tidbcloud",
    ("nextgen", "prod"): "prod-tidbcloud-nextgen",
    ("nextgen", "staging"): "staging-tidbcloud-nextgen",
    ("nextgen", "dev"): "dev-tidbcloud-nextgen",
    # byoc rides the nextgen biz type until the Clusters API grows a byoc type.
    ("byoc", "prod"): "prod-tidbcloud-nextgen",
    ("byoc", "staging"): "staging-tidbcloud-nextgen",
    ("byoc", "dev"): "dev-tidbcloud-nextgen",
    ("serverless", "prod"): "prod-devtier",
    ("serverless", "staging"): "staging-devtier",
    ("serverless", "dev"): "dev-devtier",
}

# Fixed applicationIDs for infrastructure clusters, used to narrow API results.
APPLICATION_ID_MAP = {
    "dedicated": "seed",
    "serverless": "devtier-infra",
}


def biz_type_for(architecture, environment):
    return BIZ_TYPE_MAP.get((architecture, environment))


# ── generation ──────────────────────────────────────────────────────────────


def extract_uuid(url):
    m = UUID_RE.search(url)
    return m.group(0) if m else None


def load_configured(path):
    """Load already-configured entries from an existing JSON config file.

    Returns a dict mapping regionId -> full entry for entries that have both
    a valid datasource URL and labels.
    """
    if not path.exists():
        return {}
    existing = json.loads(path.read_text())
    configured = {}
    for e in existing.get("entries", []):
        if not isinstance(e, dict):
            continue
        rid = e.get("regionId")
        url = (e.get("datasource") or {}).get("url")
        labels = e.get("labels")
        if isinstance(rid, str) and isinstance(url, str) and url and isinstance(labels, dict) and labels:
            configured[rid] = dict(regionId=rid, datasource=dict(url=url), labels=dict(labels))
    return configured


def generate(regions, router_map, primary_by_region, routes, phys_regions, configured=None):
    skipped = []
    entries = []
    for region_entry in regions.get("entries", []):
        if not isinstance(region_entry, dict):
            continue
        rid = region_entry.get("regionId")
        if not isinstance(rid, str):
            continue

        if configured and rid in configured:
            entries.append(configured[rid])
            continue

        region = parse_region_id(rid)
        if not region:
            skipped.append(dict(regionId=rid, reason="invalid-region-id"))
            continue

        candidates = resolve_candidate_bases(
            region, router_map, primary_by_region, routes, phys_regions,
        )
        if not candidates:
            skipped.append(dict(regionId=rid, reason="no-route"))
            continue

        base = candidates[0]
        url = base + VM_PATH_PREFIX

        entry = dict(
            regionId=rid,
            datasource=dict(url=url),
            _candidates=candidates,
            _needs_uuid=True,
        )
        entries.append(entry)

    entries.sort(key=lambda e: e["regionId"])
    return entries, skipped


# ── Auth0 + Clusters API ────────────────────────────────────────────────────


def get_auth0_token():
    data = urlencode(dict(
        grant_type="client_credentials",
        client_id=AUTH0_CLIENT_ID,
        client_secret=AUTH0_CLIENT_SECRET,
        audience=AUTH0_AUDIENCE,
    )).encode("utf-8")
    req = Request(AUTH0_URL, data=data, method="POST")
    req.add_header("Content-Type", "application/x-www-form-urlencoded")
    with urlopen(req, timeout=30) as r:
        body = json.loads(r.read())
    token = body.get("access_token")
    if not token:
        raise RuntimeError("No access_token in Auth0 response")
    return token


def fetch_clusters(token, vendor, region, biz_type=None, application_id=None, project_id=None):
    params = dict(vendor=vendor, region=region)
    if biz_type:
        params["bizType"] = biz_type
    if application_id:
        params["applicationID"] = application_id
    if project_id:
        params["projectID"] = project_id
    qs = urlencode(params)
    url = f"{CLUSTERS_API_BASE}?{qs}"
    req = Request(url)
    req.add_header("Authorization", f"Bearer {token}")
    try:
        with urlopen(req, timeout=30) as r:
            body = json.loads(r.read())
        return body.get("data", {}).get("clusters", [])
    except HTTPError as e:
        print(f"    HTTP {e.code} from {url}", file=sys.stderr)
        return []
    except Exception as e:
        print(f"    error fetching {url}: {e}", file=sys.stderr)
        return []


# ── enrichment ──────────────────────────────────────────────────────────────


def enrich(entries, token, phys_regions):
    # Collect unique (vendor, physical_region, bizType, applicationID, projectID) tuples.
    # For serverless, projectID = regionName allows precise per-shard matching.
    query_keys = set()
    entry_biz = {}
    entry_app = {}
    entry_project = {}
    for e in entries:
        if not e.get("_needs_uuid"):
            continue
        region = parse_region_id(e["regionId"])
        if not region:
            continue
        biz = biz_type_for(region["architecture"], region["environment"])
        app_id = APPLICATION_ID_MAP.get(region["architecture"])
        project_id = region["region_name"] if region["architecture"] == "serverless" else None
        entry_biz[e["regionId"]] = biz
        entry_app[e["regionId"]] = app_id
        entry_project[e["regionId"]] = project_id
        for rk in physical_region_keys(region["region_name"]):
            if rk in phys_regions.get(region["provider"], set()):
                query_keys.add((region["provider"], rk, biz, app_id, project_id))
                # Some seed clusters use region_name as applicationID
                # instead of the default ("seed" for dedicated, none for nextgen).
                if region["architecture"] == "dedicated":
                    if rk != app_id:
                        query_keys.add((region["provider"], rk, biz, rk, project_id))
                    if region["region_name"] != app_id and region["region_name"] != rk:
                        query_keys.add((region["provider"], rk, biz, region["region_name"], project_id))
                if region["architecture"] in ("nextgen", "byoc"):
                    query_keys.add((region["provider"], rk, biz, region["region_name"], project_id))

    print(f"  fetching clusters for {len(query_keys)} query keys")
    all_clusters = {}
    for vendor, region, biz, app_id, proj_id in sorted(query_keys, key=lambda t: (t[0], t[1], t[2] or "", t[4] or "")):
        key = (vendor, region, biz, app_id, proj_id)
        if key in all_clusters:
            continue
        clusters = fetch_clusters(token, vendor, region, biz_type=biz, application_id=app_id, project_id=proj_id)
        all_clusters[key] = clusters
        label = f"bizType={biz}" + (f" appID={app_id}" if app_id else "") + (f" projectID={proj_id}" if proj_id else "")
        print(f"    {vendor}-{region} {label}: {len(clusters)} clusters")

    # Match clusters to entries.
    unresolved = []
    resolved = 0
    for i, e in enumerate(entries):
        if not e.pop("_needs_uuid", False):
            continue

        region = parse_region_id(e["regionId"])
        if not region:
            unresolved.append(dict(regionId=e["regionId"], reason="invalid-region-id"))
            continue

        biz = entry_biz.get(e["regionId"])
        app_id = entry_app.get(e["regionId"])
        proj_id = entry_project.get(e["regionId"])
        candidates = e.get("_candidates", [])
        matched = False

        for base in candidates:
            full_prefix = base + VM_PATH_PREFIX
            for rk in physical_region_keys(region["region_name"]):
                if rk not in phys_regions.get(region["provider"], set()):
                    continue
                # Prefer specific applicationIDs over the default ("seed" / None)
                # to avoid matching the wrong cluster on shared dataplanes.
                lookup_ids = []
                if region["architecture"] == "dedicated":
                    if region["region_name"] != app_id and region["region_name"] != rk:
                        lookup_ids.append(region["region_name"])
                    if rk != app_id:
                        lookup_ids.append(rk)
                    lookup_ids.append(app_id)
                elif region["architecture"] in ("nextgen", "byoc"):
                    lookup_ids.append(region["region_name"])
                    lookup_ids.append(app_id)
                else:
                    lookup_ids.append(app_id)
                for aid in lookup_ids:
                    for c in all_clusters.get((region["provider"], rk, biz, aid, proj_id), []):
                        uri = c.get("internalReadUri", "")
                        if uri.startswith(full_prefix):
                            # Validate that the matched cluster's applicationID
                            # matches the region.  Without this check, a generic
                            # (appID=None) query on a shared dataplane can match
                            # a cluster belonging to a different application.
                            c_app = c.get("applicationID")
                            if aid is None and c_app and c_app != region["region_name"]:
                                continue
                            uri = uri.rstrip("/")
                            if uri.endswith("/api/v1"):
                                uri = uri[:-len("/api/v1")]
                            e["datasource"] = dict(url=uri)
                            matched = True
                            break
                    if matched:
                        break
                if matched:
                    break
            if matched:
                    break

        if matched:
            resolved += 1
            print(f"  resolved {e['regionId']}")
        else:
            unresolved.append(dict(regionId=e["regionId"], reason="no-cluster-match"))
            print(f"  unresolved {e['regionId']}")

    print(f"  resolved {resolved}, unresolved {len(unresolved)}")
    return entries, unresolved


# ── verification + label enrichment ────────────────────────────────────────


def query_acs_build_info(url):
    """Query acs_build_info from a VM datasource URL."""
    api = f"{url.rstrip('/')}/api/v1/query?{urlencode(dict(query='acs_build_info'))}"
    try:
        with urlopen(Request(api), timeout=15) as r:
            body = json.loads(r.read())
        rows = []
        for item in body.get("data", {}).get("result") or []:
            m = item.get("metric", {})
            rows.append(m)
        return dict(ok=True, metrics=rows, error=None, status=None)
    except HTTPError as e:
        return dict(ok=False, metrics=[], error=f"HTTP {e.code}", status=e.code)
    except Exception as e:
        return dict(ok=False, metrics=[], error=type(e).__name__, status=None)


def _choose_cluster(region, clusters):
    """Pick the best-matching cluster label from acs_build_info results.

    When a VM instance serves multiple clusters, we need to select the one
    that matches the region rather than just taking the first result.
    """
    if not clusters:
        return None
    if len(clusters) == 1:
        return clusters[0]

    name = region["region_name"]
    stripped = SUFFIX_RE.sub("", name)
    plat = {"aws": "eks", "gcp": "gke", "azure": "aks"}.get(region["provider"], "ack")

    # Exact-match candidates in priority order.
    seen = set()
    candidates = []
    for c in [
        f"seed/{plat}/{name}", f"seed/{plat}/{stripped}",
        f"freetier/{plat}/{name}", f"freetier/{plat}/{stripped}",
        name, stripped,
    ]:
        if c and c not in seen:
            seen.add(c)
            candidates.append(c)

    for c in candidates:
        if c in clusters:
            return c

    # Fallback: substring scoring.
    best, best_score = None, 0
    for cl in clusters:
        score = 0
        if name in cl:
            score += 40
        if stripped and stripped in cl:
            score += 25
        if cl.startswith("seed/") or cl.startswith("freetier/"):
            score += 10
        if score > best_score:
            best, best_score = cl, score

    return best if best_score > 0 else clusters[0]


def verify_and_label(entries):
    """Verify URLs and enrich entries with cluster labels from acs_build_info.

    Skips entries that already have labels (preserved from existing config).
    For new entries, tries the current URL first.  On 4xx errors, falls back
    to candidate base URLs from routing (like the Loki script's 4xx fallback).
    """
    # Build full candidate URLs for each entry that needs verification.
    entry_candidates = []
    all_urls = set()
    for e in entries:
        # Skip already-configured entries (have labels, no _candidates).
        if e.get("labels") and not e.get("_candidates"):
            entry_candidates.append([])
            continue

        url = e.get("datasource", {}).get("url", "")
        uuid = extract_uuid(url)
        if not uuid:
            entry_candidates.append([])
            continue

        # Current URL is always first, then routing-based alternatives.
        full_urls = [url]
        for base in e.get("_candidates", []):
            fallback = base + VM_PATH_PREFIX + uuid
            if fallback not in full_urls:
                full_urls.append(fallback)
        entry_candidates.append(full_urls)
        all_urls.update(full_urls)

    all_urls = sorted(all_urls)
    total_urls = len(all_urls)

    reports = {}
    for idx, u in enumerate(all_urls, 1):
        print(f"  querying [{idx}/{total_urls}] {u}")
        reports[u] = query_acs_build_info(u)

    total = len(entries)
    unresolved = []
    for i, e in enumerate(entries):
        candidates = entry_candidates[i]
        if not candidates:
            continue

        region = parse_region_id(e["regionId"])
        found = False
        attempt_details = []

        for url in candidates:
            report = reports.get(url, {})
            error = report.get("error")

            if error:
                attempt_details.append(dict(url=url, error=error))
                print(f"  verifying [{i+1}/{total}] {e['regionId']}")
                print(f"    {url}")
                print(f"    {error}")
                # 4xx: try next candidate; other errors: break.
                status = report.get("status")
                if not (isinstance(status, int) and 400 <= status < 500):
                    break
                continue

            metrics = report.get("metrics", [])
            if not metrics:
                attempt_details.append(dict(url=url, error="no data"))
                print(f"  verifying [{i+1}/{total}] {e['regionId']}")
                print(f"    {url}")
                print(f"    no data")
                continue

            all_clusters = list(dict.fromkeys(
                m["cluster"] for m in metrics if m.get("cluster")
            ))
            cluster = _choose_cluster(region, all_clusters) if region else (all_clusters[0] if all_clusters else None)

            if cluster:
                e["datasource"] = dict(url=url)
                e["labels"] = dict(
                    container=CONTAINER,
                    cluster=cluster,
                )
                print(f"  verifying [{i+1}/{total}] {e['regionId']} → {cluster}")
                found = True
                break
            else:
                print(f"  verifying [{i+1}/{total}] {e['regionId']}")
                print(f"    {url}")
                print(f"    no cluster label")

        if not found and candidates:
            print(f"  verifying [{i+1}/{total}] {e['regionId']} — all candidates failed")
            for a in attempt_details:
                print(f"    {a['url']}")
                print(f"      {a['error']}")
            unresolved.append(dict(
                regionId=e["regionId"],
                attempts=attempt_details,
            ))

    return entries, unresolved


# ── main ────────────────────────────────────────────────────────────────────


def main():
    print(f"loading {REGIONS_PATH}")
    regions = json.loads(REGIONS_PATH.read_text())

    if not regions.get("entries"):
        print("error: regions.json has no entries", file=sys.stderr)
        sys.exit(1)

    configured = load_configured(VM_PATH)
    if configured:
        print(f"  {len(configured)} already-configured entries from {VM_PATH.name}")

    print(f"loading {META_PATH}")
    router_map, primary_by_region, routes, phys_regions = load_meta(META_PATH)
    print(f"  {len(router_map)} routers, {len(routes)} route entries")

    entries, skipped = generate(
        regions, router_map, primary_by_region, routes, phys_regions, configured,
    )
    print(f"generated {len(entries)} entries ({len(configured)} preserved, {len(entries) - len(configured)} new), skipped {len(skipped)}")

    needs_uuid = sum(1 for e in entries if e.get("_needs_uuid"))
    if skipped:
        for s in skipped:
            print(f"  skipped: {s['regionId']} ({s['reason']})")

    # Enrich with Clusters API if Auth0 credentials are available.
    if needs_uuid and AUTH0_CLIENT_ID and AUTH0_CLIENT_SECRET:
        print("enriching cluster UUIDs via Clusters API")
        token = get_auth0_token()
        entries, unresolved = enrich(entries, token, phys_regions)
        needs_uuid = sum(1 for e in entries if e.get("_needs_uuid"))
    elif needs_uuid:
        print(f"  {needs_uuid} entries need UUID (set AUTH0_CLIENT_ID/AUTH0_CLIENT_SECRET to enrich)")

    # Drop entries without a valid cluster UUID — they are not queryable.
    before = len(entries)
    entries = [e for e in entries if extract_uuid(e.get("datasource", {}).get("url", ""))]
    dropped = before - len(entries)
    if dropped:
        print(f"  dropped {dropped} entries without UUID")

    # Verify URLs and enrich with cluster labels from acs_build_info.
    print("verifying URLs via acs_build_info")
    entries, unresolved = verify_and_label(entries)

    # Drop entries without labels — unverifiable data is unreliable.
    before = len(entries)
    entries = [e for e in entries if e.get("labels")]
    dropped = before - len(entries)
    if dropped:
        print(f"  dropped {dropped} entries without labels")

    # Strip internal fields before output.
    for e in entries:
        e.pop("_candidates", None)
        e.pop("_needs_uuid", None)

    VM_PATH.write_text(
        json.dumps(dict(version=1, entries=entries), indent=2) + "\n",
    )
    print(f"wrote {VM_PATH}")

    if unresolved:
        print(f"\nunresolved: {len(unresolved)}")
        for u in unresolved[:20]:
            print(f"  - {u['regionId']}")
            for a in u.get("attempts", []):
                print(f"      {a['url']}")
                print(f"        {a['error']}")
        sys.exit(2)

    if skipped:
        sys.exit(2)


if __name__ == "__main__":
    main()