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

Reads the existing loki.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, Loki API enrichment).  This
allows manual overrides to persist across regeneration and speeds up the
process by skipping already-resolved entries.

Routing algorithm
~~~~~~~~~~~~~~~~~
For each regionId in config/regions.json:

1. Parse regionId ``arch:provider:env:dirname`` to extract architecture,
   provider, environment, and the physical region name.

2. Construct a **biz string** from (architecture, environment):
     dedicated  + prod    → ``tidbcloud:x:x``
     dedicated  + dev     → ``dev-tidbcloud:x:x``
     dedicated  + staging → ``staging-tidbcloud:x:x``
     nextgen    + *       → ``{env}-tidbcloud-nextgen:x:x``
     byoc       + *       → ``{env}-tidbcloud-nextgen:x:x``
     serverless + *       → ``{env}-devtier:x:x``

3. **Direct logical match** — if the region name itself is a registered
   logicalRegionName in meta.json vendors (but not a physical region), use
   that dataplaneAddress directly.  Handles byoc-poc, na-test, etc.

4. **Route-based matching** — look up ``{provider}-{physical_region}`` in
   meta.json routes, match the biz string against route patterns using
   most-specific-wins (specificity scoring aligned with Go global server's
   ``calculatePatternSpecificity``).  This resolves:
     - nextgen  → ``*-premium`` or ``*-ng`` endpoints
     - dev/staging dedicated → ``*-staging`` endpoints
     - serverless dev → ``*-serverless`` endpoints

5. **Primary fallback** — always include the primary router
   (``"primary": true``) for the physical region as a fallback candidate.

6. Generate returns an ordered list of candidate URLs per entry.  During
   enrichment, each URL is tried in order; the first one returning matching
   cluster data wins.  If the routed endpoint (e.g. ``-ng``) does not have
   its own self-monitoring Loki, it falls back to the primary endpoint.

7. Final Loki URL = ``dataplaneAddress + "/loki/self-monitoring"``

Usage:
    make generate-loki
"""

import json
import math
import re
import sys
import time
from collections import defaultdict
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,
)

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"
)
OUTPUT_PATH = PROJECT / "config" / "loki.json"
LOKI_PATH = "/loki/self-monitoring"


def resolve_candidate_urls(region, router_map, primary_by_region, routes, phys_regions):
    """Return ordered candidate URLs with Loki path appended."""
    bases = resolve_candidate_bases(
        region, router_map, primary_by_region, routes, phys_regions,
    )
    if not bases:
        return None
    return [b + LOKI_PATH for b in bases]


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


def _load_configured(existing):
    """Extract fully-configured entries from existing loki.json.

    Returns a dict mapping regionId -> full entry for entries that have both
    a valid datasource URL and labels.
    """
    configured = {}
    if not existing:
        return 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(vm, router_map, primary_by_region, routes, phys_regions, existing=None):
    configured = _load_configured(existing)

    skipped = []
    entries = []
    for vm_e in vm.get("entries", []):
        if not isinstance(vm_e, dict):
            continue
        rid = vm_e.get("regionId")
        if not isinstance(rid, str):
            continue

        if 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_urls(
            region, router_map, primary_by_region, routes, phys_regions,
        )
        if not candidates:
            skipped.append(dict(regionId=rid, reason="no-route"))
            continue

        entry = dict(
            regionId=rid,
            datasource=dict(url=candidates[0]),
            _candidates=candidates,
        )
        entries.append(entry)

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


# ── Loki API query ──────────────────────────────────────────────────────────


def query_loki(url):
    try:
        query = (
            f'sum by (cluster_env,cluster) '
            f'(count_over_time({{container="{CONTAINER}"}}[14d]))'
        )
        api = (
            f"{url.rstrip('/')}/loki/api/v1/query"
            f"?{urlencode(dict(query=query, time=str(int(time.time()))))}"
        )
        with urlopen(Request(api), timeout=30) as r:
            body = json.loads(r.read())
        rows = []
        for item in body.get("data", {}).get("result") or []:
            m = item.get("metric", {})
            c = m.get("cluster")
            if not isinstance(c, str):
                continue
            ce = m.get("cluster_env")
            cnt = float((item.get("value") or [None, "0"])[1] or 0)
            rows.append(dict(
                cluster_env=ce if isinstance(ce, str) else None,
                cluster=c, count=cnt,
            ))
        return dict(ok=True, rows=rows, error=None, status=None)
    except HTTPError as e:
        return dict(ok=False, rows=[], error=str(e), status=e.code)
    except Exception as e:
        return dict(ok=False, rows=[], error=str(e), status=None)


# ── cluster resolution ──────────────────────────────────────────────────────


def platform_for(prov):
    return {"aws": "eks", "gcp": "gke", "azure": "aks"}.get(prov, "ack")


def candidate_clusters(region):
    name = region["region_name"]
    stripped = SUFFIX_RE.sub("", name)
    plat = platform_for(region["provider"])
    env = region["environment"]
    seen, out = set(), []
    for c in [
        name, stripped,
        f"{env}-{name}", f"{env}-{stripped}",
        f"seed/{plat}/{name}", f"seed/{plat}/{stripped}",
    ]:
        if c and c not in seen:
            seen.add(c)
            out.append(c)
    return out


def rank_cluster(region, cluster, count):
    stripped = SUFFIX_RE.sub("", region["region_name"])
    s = 0.0
    if region["architecture"] == "serverless":
        if (cluster.startswith(f"{region['environment']}-")
                and re.search(r"-f\d{2}$", cluster, re.I)):
            s += 70
    else:
        if cluster.startswith("seed/"):
            s += 70
    if region["region_name"] in cluster:
        s += 40
    if stripped in cluster:
        s += 25
    s += min(math.log10(count + 1), 8)
    return s


def choose_cluster(region, rows):
    env_rows = [
        r for r in rows
        if r["cluster_env"] == region["environment"] and r["cluster"]
    ]
    if not env_rows:
        all_envs = {r["cluster_env"] for r in rows if r["cluster"]}
        return None, (
            f"no rows with cluster_env={region['environment']}; "
            f"available envs: {sorted(all_envs) or 'none'}; "
            f"total rows: {len(rows)}"
        )

    cands = candidate_clusters(region)
    exact = defaultdict(list)
    for r in env_rows:
        if r["cluster"] in cands:
            exact[r["cluster"]].append(r)

    for c in cands:
        hits = exact.get(c, [])
        if hits:
            return max(hits, key=lambda h: h["count"])["cluster"], None

    ranked = sorted(
        [(r["cluster"], rank_cluster(region, r["cluster"], r["count"]))
         for r in env_rows
         if region["region_name"] in r["cluster"]],
        key=lambda x: x[1], reverse=True,
    )
    if not ranked:
        return None, "no rows after ranking"
    if len(ranked) == 1:
        return ranked[0][0], None
    if ranked[0][1] < 30:
        return None, (
            f"best score too low: {ranked[0][0]}={ranked[0][1]:.1f}; "
            f"top 3: {[(c, f'{s:.1f}') for c, s in ranked[:3]]}"
        )
    if ranked[0][1] == ranked[1][1]:
        return None, (
            f"tie: {ranked[0][0]}={ranked[0][1]:.1f} "
            f"vs {ranked[1][0]}={ranked[1][1]:.1f}"
        )
    return ranked[0][0], None


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


def enrich(entries):
    # Collect all unique candidate URLs across entries that need enrichment.
    all_urls = set()
    for e in entries:
        # Skip already-configured entries (have labels, no _candidates).
        if e.get("labels") and not e.get("_candidates"):
            continue
        for u in e.get("_candidates", []):
            all_urls.add(u)
    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_loki(u)

    unresolved = []
    total = len(entries)
    for i, e in enumerate(entries):
        # Skip already-configured entries.
        if e.get("labels") and not e.get("_candidates"):
            continue

        region = parse_region_id(e["regionId"])
        if not region:
            print(f"  resolving [{i+1}/{total}] {e['regionId']} — invalid-region-id")
            unresolved.append(
                dict(regionId=e["regionId"], reason="invalid-region-id"),
            )
            continue

        # Try each candidate URL in order; first one with cluster data wins.
        # If a candidate returns an API error, stop — don't fall back to
        # another URL which may return data for a different cluster.
        candidates = e.get("_candidates", [e["datasource"]["url"]])
        found = False
        attempt_details = []
        for url in candidates:
            report = reports.get(url, {})
            api_error = report.get("error")

            if api_error:
                attempt_details.append(dict(url=url, error=api_error))
                print(f"  resolving [{i+1}/{total}] {e['regionId']} [{url}] — api-error: {api_error}")
                status = report.get("status")
                if not (isinstance(status, int) and 400 <= status < 500):
                    break

            api_rows = report.get("rows", [])
            cluster, detail = choose_cluster(region, api_rows)

            if cluster:
                print(f"  resolving [{i+1}/{total}] {e['regionId']} [{url}] → {cluster}")
                entries[i] = dict(
                    regionId=e["regionId"],
                    datasource=dict(url=url),
                    labels=dict(
                        cluster_env=region["environment"],
                        container=CONTAINER,
                        cluster=cluster,
                    ),
                )
                found = True
                break
            else:
                row_count = len(api_rows)
                envs = sorted({r["cluster_env"] for r in api_rows if r.get("cluster_env")})
                attempt_details.append(dict(
                    url=url, rows=row_count,
                    envs=envs or None, detail=detail,
                ))
                print(f"  resolving [{i+1}/{total}] {e['regionId']} [{url}] — {detail}")

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

    # 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 _candidates from entries before output.
    for e in entries:
        e.pop("_candidates", None)

    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)

    existing = None
    if OUTPUT_PATH.exists():
        print(f"loading {OUTPUT_PATH}")
        existing = json.loads(OUTPUT_PATH.read_text())

    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, existing,
    )
    configured_count = sum(1 for e in entries if not e.get("_candidates"))
    print(f"generated {len(entries)} entries ({configured_count} preserved, {len(entries) - configured_count} new), skipped {len(skipped)}")

    if skipped:
        for s in skipped:
            print(f"  skipped: {s['regionId']} ({s['reason']})")

    print("enriching labels from Loki APIs")
    entries, unresolved = enrich(entries)

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

    if unresolved:
        print(f"unresolved: {len(unresolved)}")
        for u in unresolved[:20]:
            print(f"  - {u['regionId']}")
            for a in u.get("attempts", []):
                reason = a.get("detail") or a.get("error", "")
                envs = a.get("envs")
                print(f"      {a['url']}")
                print(f"        {reason}")
                if envs:
                    print(f"        available envs: {envs}")
        sys.exit(2)


if __name__ == "__main__":
    main()