#!/usr/bin/env python3
"""
Write a CSV of Power BI semantic models and the reports (and dependent
semantic models) that sit downstream of them, using the Power BI admin
Scanner API.

Flow (all calls use the Power BI REST API, v1.0, "myorg"):
  1. GET  /admin/workspaces/modified               list workspace IDs
  2. POST /admin/workspaces/getInfo                start a scan (up to 100 workspaces)
  3. GET  /admin/workspaces/scanStatus/{scanId}    poll until the scan has succeeded
  4. GET  /admin/workspaces/scanResult/{scanId}    download the result

Authentication: a Microsoft Entra app (service principal) using the client
credentials flow. Secrets are read from environment variables, never from
the command line or from this file.

  PBI_TENANT_ID       Entra tenant ID (GUID) or domain name
  PBI_CLIENT_ID       application (client) ID of the app registration
  PBI_CLIENT_SECRET   client secret for that app
  PBI_WORKSPACE_IDS   optional, comma separated workspace IDs. Use it for a
                      small first run instead of scanning every workspace.

Prerequisites (a Fabric administrator sets these up, see Microsoft Learn,
"Set up metadata scanning" and "Enable service principal authentication
for admin APIs"):
  * the app is allowed to use read-only admin APIs
  * the app has no admin-consent-required Power BI permissions in Entra
  * for --measures, both "Enhance admin APIs responses with detailed
    metadata" and "Enhance admin APIs responses with DAX and mashup
    expressions" are enabled

Documented limits at the time of writing (check the Learn pages again before
a large run): getInfo 500 requests per hour, 16 simultaneous requests,
1 to 100 workspaces per call; scanResult 500 requests per hour and the
result stays available for 24 hours; scanStatus 10,000 requests per hour;
modified 30 requests per hour. This script runs one scan at a time.

Only the Python standard library is used.

Usage:
  python3 scan_semantic_model_dependents.py --out dependents.csv
  python3 scan_semantic_model_dependents.py --out dependents.csv --measures measures.csv
"""

import argparse
import csv
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

API = "https://api.powerbi.com/v1.0/myorg"
TOKEN_URL = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
SCOPE = "https://analysis.windows.net/powerbi/api/.default"
BATCH = 100  # documented maximum workspaces per getInfo call


def need_env(name):
    value = os.environ.get(name)
    if not value:
        sys.exit("Missing environment variable: " + name)
    return value


def get_token():
    body = urllib.parse.urlencode({
        "client_id": need_env("PBI_CLIENT_ID"),
        "client_secret": need_env("PBI_CLIENT_SECRET"),
        "scope": SCOPE,
        "grant_type": "client_credentials",
    }).encode()
    req = urllib.request.Request(
        TOKEN_URL.format(tenant=need_env("PBI_TENANT_ID")),
        data=body,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
    )
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            return json.load(resp)["access_token"]
    except urllib.error.HTTPError as err:
        detail = err.read().decode(errors="replace")[:500]
        sys.exit("Token request failed with HTTP %d\n%s" % (err.code, detail))


def call(method, url, token, payload=None):
    """One API call. Waits and retries on HTTP 429 using Retry-After."""
    data = json.dumps(payload).encode() if payload is not None else None
    for attempt in range(6):
        req = urllib.request.Request(url, data=data, method=method)
        req.add_header("Authorization", "Bearer " + token)
        if data is not None:
            req.add_header("Content-Type", "application/json")
        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                raw = resp.read()
                return json.loads(raw) if raw else {}
        except urllib.error.HTTPError as err:
            if err.code == 429 and attempt < 5:
                wait = int(err.headers.get("Retry-After", "60"))
                print("Throttled (429). Waiting %d seconds." % wait, file=sys.stderr)
                time.sleep(wait)
                continue
            detail = err.read().decode(errors="replace")[:500]
            sys.exit("HTTP %d from %s %s\n%s" % (err.code, method, url, detail))
    sys.exit("Gave up after repeated throttling: " + url)


def list_workspace_ids(token):
    url = API + "/admin/workspaces/modified?excludePersonalWorkspaces=true"
    return [w["id"] for w in call("GET", url, token)]


def scan(token, workspace_ids, with_measures):
    query = {"lineage": "true", "datasourceDetails": "true"}
    if with_measures:
        query["datasetSchema"] = "true"
        query["datasetExpressions"] = "true"
    url = API + "/admin/workspaces/getInfo?" + urllib.parse.urlencode(query)
    started = call("POST", url, token, {"workspaces": workspace_ids})
    scan_id = started["id"]

    status_url = API + "/admin/workspaces/scanStatus/" + scan_id
    for _ in range(120):  # up to about 20 minutes
        status = call("GET", status_url, token)
        state = status.get("status")
        if state == "Succeeded":
            break
        if status.get("error"):
            sys.exit("Scan %s reported an error: %s" % (scan_id, status["error"]))
        print("Scan %s status: %s" % (scan_id, state), file=sys.stderr)
        time.sleep(10)
    else:
        sys.exit("Scan %s did not report Succeeded in time." % scan_id)

    return call("GET", API + "/admin/workspaces/scanResult/" + scan_id, token)


def build_rows(workspaces):
    """One row per (semantic model, dependent). Models with none get one row."""
    models = {}  # dataset id -> (workspace dict, dataset dict)
    for ws in workspaces:
        for ds in ws.get("datasets", []):
            models[ds["id"]] = (ws, ds)

    dependents = {}  # dataset id -> list of dependent rows

    for ws in workspaces:
        for rep in ws.get("reports", []):
            ds_id = rep.get("datasetId")
            if not ds_id:
                continue
            # datasetWorkspaceId is returned only when the model lives in a
            # different workspace from the report.
            model_ws_id = rep.get("datasetWorkspaceId") or ws["id"]
            dependents.setdefault(ds_id, []).append({
                "dependent_type": "report",
                "dependent_workspace_id": ws["id"],
                "dependent_workspace": ws.get("name", ""),
                "dependent_id": rep.get("id", ""),
                "dependent_name": rep.get("name", ""),
                "report_type": rep.get("reportType", ""),
                "modified_by": rep.get("modifiedBy", ""),
                "modified_date": rep.get("modifiedDateTime", ""),
                "cross_workspace": str(model_ws_id != ws["id"]).lower(),
            })

        for ds in ws.get("datasets", []):
            for up in ds.get("upstreamDatasets", []) or []:
                target = up.get("targetDatasetId")
                if not target:
                    continue
                dependents.setdefault(target, []).append({
                    "dependent_type": "semantic model (upstream link)",
                    "dependent_workspace_id": ws["id"],
                    "dependent_workspace": ws.get("name", ""),
                    "dependent_id": ds.get("id", ""),
                    "dependent_name": ds.get("name", ""),
                    "report_type": "",
                    "modified_by": "",
                    "modified_date": "",
                    "cross_workspace": str(up.get("groupId", ws["id"]) != ws["id"]).lower(),
                })

    blank = {
        "dependent_type": "", "dependent_workspace_id": "", "dependent_workspace": "",
        "dependent_id": "", "dependent_name": "", "report_type": "",
        "modified_by": "", "modified_date": "", "cross_workspace": "",
    }

    rows = []
    # Models that appear only as targets (their workspace was not scanned)
    # still get rows, with empty model details.
    for ds_id in sorted(set(models) | set(dependents)):
        ws, ds = models.get(ds_id, ({}, {}))
        endorsement = (ds.get("endorsementDetails") or {}).get("endorsement", "")
        head = {
            "workspace_id": ws.get("id", ""),
            "workspace": ws.get("name", ""),
            "semantic_model_id": ds_id,
            "semantic_model": ds.get("name", ""),
            "configured_by": ds.get("configuredBy", ""),
            "endorsement": endorsement,
        }
        for dep in dependents.get(ds_id) or [blank]:
            rows.append({**head, **dep})
    return rows


def write_csv(path, rows):
    fields = [
        "workspace_id", "workspace", "semantic_model_id", "semantic_model",
        "configured_by", "endorsement", "dependent_type",
        "dependent_workspace_id", "dependent_workspace", "dependent_id",
        "dependent_name", "report_type", "modified_by", "modified_date",
        "cross_workspace",
    ]
    with open(path, "w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def write_measures(path, workspaces):
    fields = ["workspace", "semantic_model_id", "semantic_model", "table",
              "measure", "expression", "is_hidden"]
    with open(path, "w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        for ws in workspaces:
            for ds in ws.get("datasets", []):
                for table in ds.get("tables", []) or []:
                    for m in table.get("measures", []) or []:
                        writer.writerow({
                            "workspace": ws.get("name", ""),
                            "semantic_model_id": ds.get("id", ""),
                            "semantic_model": ds.get("name", ""),
                            "table": table.get("name", ""),
                            "measure": m.get("name", ""),
                            "expression": m.get("expression", ""),
                            "is_hidden": m.get("isHidden", ""),
                        })


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--out", required=True, help="CSV file for semantic models and dependents")
    parser.add_argument("--measures", help="optional CSV file for measures and their DAX")
    args = parser.parse_args()

    token = get_token()
    only = os.environ.get("PBI_WORKSPACE_IDS", "").replace(" ", "")
    ids = only.split(",") if only else list_workspace_ids(token)
    print("Workspaces to scan: %d" % len(ids), file=sys.stderr)

    workspaces = []
    for i in range(0, len(ids), BATCH):
        chunk = ids[i:i + BATCH]
        print("Scanning workspaces %d to %d" % (i + 1, i + len(chunk)), file=sys.stderr)
        result = scan(token, chunk, with_measures=bool(args.measures))
        workspaces.extend(result.get("workspaces", []))

    rows = build_rows(workspaces)
    write_csv(args.out, rows)
    print("Wrote %d rows to %s" % (len(rows), args.out), file=sys.stderr)

    if args.measures:
        write_measures(args.measures, workspaces)
        print("Wrote measures to %s" % args.measures, file=sys.stderr)


if __name__ == "__main__":
    main()
