#!/usr/bin/env python3
"""Bulk-add a Power BI service principal to workspaces using OAuth device code and
the Admin AddUserAsAdmin API. Standard library only (urllib)."""

import argparse
import json
import os
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from typing import Any

API_ROOT = "https://api.powerbi.com/v1.0/myorg"
SCOPE = "https://analysis.windows.net/powerbi/api/Tenant.ReadWrite.All offline_access"


def configure_stdio() -> None:
    """Use line-buffered stdout/stderr when not a TTY so device-code prompts show up immediately."""
    for stream in (sys.stdout, sys.stderr):
        try:
            if hasattr(stream, "reconfigure"):
                stream.reconfigure(line_buffering=True)
        except (OSError, ValueError, AttributeError):
            pass


def die(msg: str, code: int = 1) -> None:
    print(f"error: {msg}", file=sys.stderr, flush=True)
    raise SystemExit(code)


def http_request(
    url: str,
    *,
    method: str = "GET",
    headers: dict[str, str] | None = None,
    body: bytes | None = None,
    timeout: int = 120,
) -> tuple[int, bytes]:
    h = dict(headers or {})
    req = urllib.request.Request(url, data=body, method=method, headers=h)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.getcode(), resp.read()
    except urllib.error.HTTPError as e:
        return e.code, e.read()


def post_urlencoded(
    url: str,
    fields: dict[str, str],
    *,
    timeout: int = 60,
    ok_http_error: frozenset[int] | None = None,
) -> dict[str, Any]:
    """POST application/x-www-form-urlencoded; parse JSON body.

    Token endpoint may return HTTP 400 with {"error":"authorization_pending"}.
    """
    data = urllib.parse.urlencode(fields).encode()
    code, raw = http_request(
        url,
        method="POST",
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        body=data,
        timeout=timeout,
    )
    try:
        out: dict[str, Any] = json.loads(raw.decode())
    except json.JSONDecodeError:
        out = {}
    if code < 400:
        return out
    if ok_http_error and code in ok_http_error and isinstance(out, dict):
        return out
    die(f"HTTP {code} from {url}: {raw[:500]!r}")


def obtain_access_token(
    tenant: str,
    bootstrap_client_id: str,
    *,
    access_token_file: str | None,
    poll_interval: int,
    device_verify_url: str,
) -> str:
    if access_token_file:
        path = os.path.expanduser(access_token_file)
        if not os.path.isfile(path):
            die(f"access token file not found: {access_token_file}")
        with open(path, encoding="utf-8") as f:
            token = f.readline().strip("\r\n")
        if not token:
            die("empty access token file")
        print(f"Using access token from: {access_token_file}", flush=True)
        return token

    dc_url = f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/devicecode"
    dc = post_urlencoded(
        dc_url,
        {"client_id": bootstrap_client_id, "scope": SCOPE},
        timeout=60,
    )
    if "device_code" not in dc:
        die(f"device code response parse failed: {dc}")
    device_code = dc["device_code"]
    dc_interval = int(dc.get("interval", 5))
    expires_in = int(dc["expires_in"])
    user_code = dc["user_code"]
    verification_uri = dc.get("verification_uri") or ""

    print("", flush=True)
    print("=== Device sign-in required ===", flush=True)
    print("Open this URL in a browser (Fabric / Power BI admin user):", flush=True)
    print(f"  {device_verify_url}", flush=True)
    print(f"If that page asks for a code, use:  {user_code}", flush=True)
    if verification_uri:
        print(f"(Microsoft may also show: {verification_uri})", flush=True)
    print("", flush=True)
    sys.stdout.flush()
    sys.stderr.flush()
    try:
        webbrowser.open(device_verify_url)
    except Exception:
        pass

    sleep_for = max(poll_interval, dc_interval)
    token_url = f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
    end_ts = time.time() + expires_in - 30

    while True:
        if time.time() > end_ts:
            die("device code flow timed out before completion")
        tr = post_urlencoded(
            token_url,
            {
                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                "client_id": bootstrap_client_id,
                "device_code": device_code,
            },
            timeout=60,
            ok_http_error=frozenset({400}),
        )
        err = tr.get("error") or ""
        if err == "authorization_pending":
            print(
                f"... waiting for browser sign-in (polling every {sleep_for}s)",
                flush=True,
            )
            time.sleep(sleep_for)
            continue
        if err == "slow_down":
            sleep_for += 5
            print(
                f"... slow_down from token endpoint; backing off to {sleep_for}s",
                flush=True,
            )
            time.sleep(sleep_for)
            continue
        if err in ("expired_token", "invalid_grant"):
            die(f"token endpoint error: {tr}")
        access = tr.get("access_token")
        if not access:
            die(f"failed to obtain access_token: {tr}")
        print("Signed in. Access token acquired.", flush=True)
        return str(access)


def fetch_workspaces_inventory(access_token: str) -> list[dict[str, Any]]:
    base = API_ROOT
    url: str | None = f"{base}/admin/groups?$top=5000"
    out: list[dict[str, Any]] = []
    while url:
        code, raw = http_request(
            url,
            method="GET",
            headers={"Authorization": f"Bearer {access_token}"},
            timeout=120,
        )
        if code >= 400:
            die(f"admin/groups failed HTTP {code}: {raw[:800]!r}")
        body = json.loads(raw.decode())
        for g in body.get("value", []):
            if g.get("type") == "Workspace":
                out.append({"id": g.get("id"), "name": g.get("name") or ""})
        url = body.get("@odata.nextLink") or None
    return out


def filter_workspaces(
    all_ws: list[dict[str, Any]],
    *,
    mode: str,
    workspace_ids_csv: str,
    name_contains: str,
    name_regex: str,
    skip: int,
    limit: int | None,
) -> list[dict[str, Any]]:
    picked: list[dict[str, Any]] = []
    if mode == "ids":
        want = {x.strip().lower() for x in workspace_ids_csv.split(",") if x.strip()}
        for w in all_ws:
            if str(w["id"]).lower() in want:
                picked.append(w)
    elif mode == "contains":
        s = name_contains.lower()
        for w in all_ws:
            if s in (w.get("name") or "").lower():
                picked.append(w)
    elif mode == "regex":
        r = re.compile(name_regex)
        for w in all_ws:
            if r.search(w.get("name") or ""):
                picked.append(w)
    elif mode == "all":
        picked = list(all_ws)
    else:
        die(f"bad mode: {mode}")
    if skip:
        picked = picked[skip:]
    if limit is not None:
        picked = picked[:limit]
    return picked


def post_add_user(
    access_token: str,
    workspace_id: str,
    workspace_name: str,
    principal_oid: str,
    role: str,
) -> None:
    payload = json.dumps(
        {
            "identifier": principal_oid,
            "groupUserAccessRight": role,
            "principalType": "App",
        }
    ).encode()
    url = f"{API_ROOT}/admin/groups/{workspace_id}/users"
    code, raw = http_request(
        url,
        method="POST",
        headers={
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
        },
        body=payload,
        timeout=120,
    )
    print(f"POST AddUserAsAdmin: {workspace_name} ({workspace_id})", flush=True)
    if code not in (200, 204):
        snippet = raw.decode(errors="replace").replace("\n", "")[:500]
        print(f"  HTTP {code} — {snippet}", flush=True)
    else:
        print(f"  HTTP {code} OK", flush=True)


def workspace_ids_targets(csv: str, skip: int, limit: int | None) -> list[tuple[str, str]]:
    rows: list[tuple[str, str]] = []
    for raw in csv.split(","):
        wid = "".join(raw.split())
        if wid:
            rows.append((wid, "(workspace id)"))
    if skip:
        rows = rows[skip:]
    if limit is not None:
        rows = rows[:limit]
    return rows


def parse_args(argv: list[str] | None) -> argparse.Namespace:
    env_poll = os.environ.get("POLL_INTERVAL", "5")
    try:
        default_poll = int(env_poll)
    except ValueError:
        default_poll = 5

    p = argparse.ArgumentParser(
        description="Add a Power BI app's service principal to workspaces (bulk)."
    )
    p.add_argument("--tenant", required=True, help="Microsoft Entra tenant ID")
    p.add_argument(
        "--bootstrap-client-id",
        required=True,
        help="Bootstrap app registration (public client)",
    )
    p.add_argument(
        "--principal-object-id",
        required=True,
        help="Object ID of the main app's service principal (enterprise application)",
    )
    sel = p.add_mutually_exclusive_group(required=True)
    sel.add_argument(
        "--workspace-ids",
        metavar="ID[,ID...]",
        help="Comma-separated workspace (group) UUIDs",
    )
    sel.add_argument(
        "--name-contains",
        metavar="SUBSTRING",
        help="Workspaces whose name contains this string (case-insensitive)",
    )
    sel.add_argument(
        "--name-regex",
        metavar="REGEX",
        help="Workspaces whose name matches this Python regex (re.search)",
    )
    sel.add_argument(
        "--all",
        action="store_true",
        help='All workspaces of type "Workspace" from admin inventory',
    )
    p.add_argument(
        "--skip",
        type=int,
        metavar="N",
        help="After filtering, skip the first N workspaces (then apply --limit)",
    )
    p.add_argument(
        "--limit",
        type=int,
        metavar="N",
        help="After filtering and --skip, process at most N workspaces",
    )
    p.add_argument(
        "--role",
        default="Admin",
        help="groupUserAccessRight (default: Admin)",
    )
    p.add_argument(
        "--dry-run",
        action="store_true",
        help="List targets only; no POST",
    )
    p.add_argument(
        "--poll-interval",
        type=int,
        default=default_poll,
        help=f"Device-code poll interval in seconds (default: {default_poll}; env POLL_INTERVAL)",
    )
    p.add_argument(
        "--access-token-file",
        metavar="PATH",
        help="Skip device login; read bearer token from first line",
    )
    return p.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    configure_stdio()
    args = parse_args(argv)
    device_verify_url = os.environ.get(
        "DEVICE_VERIFY_URL", "https://login.microsoft.com/device"
    )

    if args.poll_interval < 1:
        die("--poll-interval must be a positive integer")
    skip = args.skip or 0
    if skip < 0:
        die("--skip must be a non-negative integer")
    limit = args.limit
    if limit is not None and limit < 0:
        die("--limit must be a non-negative integer")

    if args.workspace_ids:
        mode = "ids"
    elif args.name_contains:
        mode = "contains"
    elif args.name_regex:
        mode = "regex"
    else:
        mode = "all"

    targets: list[tuple[str, str]]

    if mode == "ids":
        targets = workspace_ids_targets(args.workspace_ids, skip, limit)
        if args.dry_run:
            print(
                f"Matched {len(targets)} workspace id(s) (dry run; no admin inventory or sign-in):"
            )
            for wid, wname in targets:
                print(f"  {wid}\t{wname}")
            print("Dry run: no POST performed.")
            return
        token = obtain_access_token(
            args.tenant,
            args.bootstrap_client_id,
            access_token_file=args.access_token_file,
            poll_interval=args.poll_interval,
            device_verify_url=device_verify_url,
        )
        print("")
        print(f"Will POST AddUserAsAdmin for {len(targets)} workspace(s).")
        for wid, wname in targets:
            post_add_user(token, wid, wname, args.principal_object_id, args.role)
        print("")
        print("Done.")
        return

    if args.dry_run and not args.access_token_file:
        die(
            "--dry-run with --name-contains, --name-regex, or --all requires "
            "--access-token-file (or use --workspace-ids for a tokenless dry-run)"
        )

    token = obtain_access_token(
        args.tenant,
        args.bootstrap_client_id,
        access_token_file=args.access_token_file,
        poll_interval=args.poll_interval,
        device_verify_url=device_verify_url,
    )
    all_ws = fetch_workspaces_inventory(token)
    filtered = filter_workspaces(
        all_ws,
        mode=mode,
        workspace_ids_csv=args.workspace_ids or "",
        name_contains=args.name_contains or "",
        name_regex=args.name_regex or "",
        skip=skip,
        limit=limit,
    )
    targets = [(str(w["id"]), str(w.get("name") or "")) for w in filtered]

    if args.dry_run:
        print(f"Matched {len(targets)} workspace(s) (dry run):")
        for wid, wname in targets:
            print(f"  {wid}\t{wname}")
        print("Dry run: no POST performed.")
        return

    if not targets:
        print("No workspaces matched the filter.")
        return

    print("")
    print(f"Matched {len(targets)} workspace(s):")
    for wid, wname in targets:
        print(f"  {wid}\t{wname}")
    print("")

    for wid, wname in targets:
        post_add_user(token, wid, wname, args.principal_object_id, args.role)

    print("")
    print("Done.")


if __name__ == "__main__":
    main()