#!/usr/bin/env python3
"""
Provision RBAC groups, service account roles, fine-grained claims, and auth flows.

Idempotent — safe to re-run. Provisions:
  1. Global admin groups (homelab-admins)
  2. Fine-grained service/bucket/project groups (minio-*, poimen-*, paperless-*, grafana-*, sqs-*)
  3. Service account roles with custom claims (paperless-ai-agent, portfolio-agent, etc.)
  4. JWT scope mappings for fine-grained claims (minio_buckets, paperless_doctypes, etc.)
  5. OAuth2 providers with scopes (api-gw, minio, poimen, paperless, grafana)
  6. Auth flows (password grant on api-gw provider)

Usage:
    source ~/.env
    export AUTHENTIK_BOOTSTRAP_TOKEN=$(kubectl -n iam get secret authentik-secrets \
        -o jsonpath='{.data.AUTHENTIK_BOOTSTRAP_TOKEN}' | base64 -d)
    python3 scripts/iam/provision-rbac.py

DO NOT commit this file to git — .gitignore covers scripts/iam/*.py.
"""
import json
import os
import sys
import urllib.error
import urllib.request
from typing import Dict, List, Any
from pathlib import Path

# Load ~/.env for OAuth2 provider secrets
env_file = Path.home() / ".env"
if env_file.exists():
    with open(env_file) as f:
        for line in f:
            line = line.strip()
            if line.startswith("export ") and "=" in line:
                key, _, value = line[7:].partition("=")
                key = key.strip()
                value = value.strip().strip('"').strip("'")
                os.environ[key] = value

AUTHENTIK_URL = "https://authentik.riotpiao.com"
TOKEN = os.environ.get("AUTHENTIK_BOOTSTRAP_TOKEN")
if not TOKEN:
    print("Error: AUTHENTIK_BOOTSTRAP_TOKEN not set")
    print("  source ~/.env")
    print("  export AUTHENTIK_BOOTSTRAP_TOKEN=$(kubectl -n iam get secret authentik-secrets \\")
    print("      -o jsonpath='{.data.AUTHENTIK_BOOTSTRAP_TOKEN}' | base64 -d)")
    sys.exit(1)


def api(method, path, data=None):
    url = f"{AUTHENTIK_URL}{path}"
    body = json.dumps(data).encode() if data is not None else None
    req = urllib.request.Request(url, data=body, method=method, headers={
        "Authorization": f"Bearer {TOKEN}",
        "Content-Type": "application/json",
    })
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            raw = resp.read()
            return resp.status, json.loads(raw) if raw else {}
    except urllib.error.HTTPError as e:
        raw = e.read()
        try:
            parsed = json.loads(raw) if raw else {}
        except json.JSONDecodeError:
            parsed = {"raw": raw.decode(errors="replace")}
        return e.code, parsed


def die(msg):
    print(f"FATAL: {msg}", file=sys.stderr)
    sys.exit(1)


# ===========================================================================
# Group Definitions (DRY: single source of truth)
# ===========================================================================
GROUPS: Dict[str, Dict[str, Any]] = {
    "homelab-admins": {
        "description": "Cluster administrators with full access",
        "is_superuser": True,
    },
    "minio-admins": {"description": "MinIO administrators", "is_superuser": False, "minio_buckets": ["*"]},
    "minio-photos": {"description": "Photos bucket (Immich) access", "is_superuser": False, "minio_buckets": ["immich"]},
    "minio-documents": {"description": "Documents bucket (Paperless) access", "is_superuser": False, "minio_buckets": ["paperless"]},
    "minio-backups": {"description": "Backups bucket read-only access", "is_superuser": False, "minio_buckets": ["backups"]},
    "poimen-admins": {"description": "Poimen memory administrators", "is_superuser": False, "memory_projects": ["*"], "memory_visibility": "private"},
    "poimen-devs": {"description": "Dev and staging projects access", "is_superuser": False, "memory_projects": ["dev", "staging"], "memory_visibility": "internal"},
    "poimen-prod-readonly": {"description": "Production projects read-only access", "is_superuser": False, "memory_projects": ["prod"], "memory_visibility": "public"},
    "paperless-admins": {"description": "Paperless administrators", "is_superuser": False, "paperless_doctypes": ["*"]},
    "paperless-finance": {"description": "Finance documents", "is_superuser": False, "paperless_doctypes": ["invoices", "receipts", "expenses"]},
    "paperless-legal": {"description": "Legal documents", "is_superuser": False, "paperless_doctypes": ["contracts", "licenses", "agreements"]},
    "paperless-hr": {"description": "HR documents", "is_superuser": False, "paperless_doctypes": ["employment", "benefits", "payroll"]},
    "grafana-admins": {"description": "Grafana administrators", "is_superuser": False, "grafana_org_role": "Admin"},
    "grafana-editors": {"description": "Grafana dashboard editors", "is_superuser": False, "grafana_org_role": "Editor"},
    "grafana-viewers": {"description": "Grafana dashboard viewers", "is_superuser": False, "grafana_org_role": "Viewer"},
    "sqs-users": {"description": "SQS/Temporal queue read access", "is_superuser": False, "sqs_queues": ["default"]},
    "sqs-writers": {"description": "SQS/Temporal queue read/write access", "is_superuser": False, "sqs_queues": ["*"]},
    "s3-users": {"description": "S3 read access", "is_superuser": False},
    "s3-writers": {"description": "S3 read/write access", "is_superuser": False},
}

SERVICE_ACCOUNTS: Dict[str, Dict[str, Any]] = {
    "paperless-ai-agent": {
        "description": "Paperless AI plugin (auto-tagging, entity extraction)",
        "roles": ["llm:inference", "memory:write", "paperless:admin"],
        "claims": {
            "minio_buckets": ["paperless"],
            "paperless_doctypes": ["*"],
            "memory_projects": ["*"],
            "authorized_models": ["reasoning", "qwen2.5:3b"],
        },
    },
    "portfolio-agent": {
        "description": "Portfolio service agent",
        "roles": ["llm:inference", "memory:read"],
        "claims": {
            "memory_projects": ["homelab", "portfolio"],
            "memory_visibility": "public",
            "authorized_models": ["ornith:35b"],
            "minio_buckets": ["backups"],
        },
    },
    "memory-agent": {
        "description": "Memory service agent",
        "roles": ["llm:inference", "memory:read", "memory:write"],
        "claims": {
            "memory_projects": ["*"],
            "memory_visibility": "private",
            "authorized_models": ["reasoning", "ornith:35b", "qwen2.5:3b"],
        },
    },
    "temporal-worker-agent": {
        "description": "Temporal workflow worker",
        "roles": ["llm:inference", "workflow:execute", "memory:read", "memory:write"],
        "claims": {
            "memory_projects": ["*"],
            "authorized_models": ["reasoning", "ornith:35b", "qwen2.5:3b"],
        },
    },
}

SCOPE_MAPPINGS: Dict[str, Dict[str, str]] = {
    "roles": {"expression": 'return user.attributes.get("roles", [])'},
    "permissions": {"expression": 'return ["*"] if any(user.groups.filter(is_superuser=True)) else list(user.groups.values_list("name", flat=True))'},
    "minio_buckets": {"expression": 'return user.attributes.get("minio_buckets", [])'},
    "paperless_doctypes": {"expression": 'return user.attributes.get("paperless_doctypes", [])'},
    "memory_projects": {"expression": 'return user.attributes.get("memory_projects", [])'},
    "memory_visibility": {"expression": 'return user.attributes.get("memory_visibility", "public")'},
    "authorized_models": {"expression": 'return user.attributes.get("authorized_models", [])'},
    "sqs_queues": {"expression": 'return user.attributes.get("sqs_queues", [])'},
    "grafana_org_role": {"expression": 'return user.attributes.get("grafana_org_role", "Viewer")'},
}

# ===========================================================================
# Phase 1: Create/sync all groups
# ===========================================================================
print("[1/6] Ensuring groups exist...")

status, res = api("GET", "/api/v3/core/groups/?page_size=100")
if status != 200:
    die(f"GET groups -> {status} {res}")
existing_groups = {g["name"]: g for g in res["results"]}

created_count = 0
for group_name, group_spec in GROUPS.items():
    if group_name in existing_groups:
        print(f"  {group_name}: already exists")
    else:
        status, res = api("POST", "/api/v3/core/groups/", {
            "name": group_name,
            "is_superuser": group_spec.get("is_superuser", False),
        })
        if status in (200, 201):
            print(f"  {group_name}: created")
            created_count += 1
        else:
            print(f"  {group_name}: FAILED {status} {res}")

print(f"  Total: {len(GROUPS)} groups, {created_count} new")

# ===========================================================================
# Phase 2: Create/sync service account users with custom claims
# ===========================================================================
print("\n[2/6] Creating/updating service account users...")

status, res = api("GET", "/api/v3/core/users/?page_size=100")
if status != 200:
    die(f"GET users -> {status} {res}")
existing_users = {u["username"]: u for u in res["results"]}

service_pwd = os.environ.get("AUTHENTIK_SERVICE_ACCOUNT_PASSWORD", "DefaultPassword123!")

for agent_name, agent_spec in SERVICE_ACCOUNTS.items():
    if agent_name in existing_users:
        user = existing_users[agent_name]
        attrs = user.get("attributes", {})
        attrs.update(agent_spec.get("claims", {}))
        attrs["roles"] = agent_spec.get("roles", [])
        status, res = api("PATCH", f"/api/v3/core/users/{user['pk']}/", {"attributes": attrs})
        if status in (200, 201):
            print(f"  {agent_name}: claims updated")
        else:
            print(f"  {agent_name}: FAILED {status} {res}")
    else:
        status, res = api("POST", "/api/v3/core/users/", {
            "username": agent_name,
            "name": agent_spec.get("description", agent_name),
            "email": f"{agent_name}@homelab.local",
            "is_active": True,
            "is_superuser": False,
            "password": service_pwd,
            "attributes": {
                **agent_spec.get("claims", {}),
                "roles": agent_spec.get("roles", []),
            },
        })
        if status in (200, 201):
            print(f"  {agent_name}: created")
        else:
            print(f"  {agent_name}: FAILED {status} {res}")

# ===========================================================================
# Phase 3: Create scope mappings for fine-grained claims
# ===========================================================================
print("\n[3/6] Creating scope mappings for fine-grained claims...")

status, res = api("GET", "/api/v3/propertymappings/provider/scope/?page_size=100")
if status != 200:
    die(f"GET scope mappings -> {status} {res}")
existing_scopes = {m["scope_name"]: m for m in res["results"]}

for scope_name, scope_spec in SCOPE_MAPPINGS.items():
    if scope_name in existing_scopes:
        print(f"  {scope_name}: already exists")
    else:
        status, res = api("POST", "/api/v3/propertymappings/provider/scope/", {
            "name": scope_name,
            "scope_name": scope_name,
            "expression": scope_spec["expression"],
        })
        if status in (200, 201):
            print(f"  {scope_name}: created")
        else:
            print(f"  {scope_name}: FAILED {status} {res}")

# ===========================================================================
# Phase 4: Get flow UUIDs (needed for providers)
# ===========================================================================
print("\n[4/6] Fetching flow UUIDs...")

status, res = api("GET", "/api/v3/flows/instances/?page_size=100")
if status != 200:
    die(f"GET flows -> {status} {res}")

flows = {f["slug"]: f["pk"] for f in res.get("results", [])}
auth_flow = flows.get("default-provider-authorization-implicit-consent")
inval_flow = flows.get("default-provider-invalidation-flow")

if not auth_flow or not inval_flow:
    die(f"Required flows not found. auth_flow={auth_flow}, inval_flow={inval_flow}")

print(f"  authorization_flow: {auth_flow}")
print(f"  invalidation_flow: {inval_flow}")

# ===========================================================================
# Phase 5: Create OAuth2 providers with scopes
# ===========================================================================
print("\n[5/6] Creating OAuth2 providers...")

status, res = api("GET", "/api/v3/providers/oauth2/?page_size=100")
if status != 200:
    die(f"GET providers -> {status} {res}")
existing_providers = {p["name"]: p for p in res.get("results", [])}

# Fetch scope mapping PKs
status, scopes_res = api("GET", "/api/v3/propertymappings/provider/scope/?page_size=100")
if status != 200:
    print("  WARNING: could not fetch scope mappings")
    scope_pks = {}
else:
    scope_pks = {m["scope_name"]: m["pk"] for m in scopes_res.get("results", [])}

scope_pks_list = [scope_pks[s] for s in SCOPE_MAPPINGS.keys() if s in scope_pks]

OAuth2_PROVIDERS = {
    "api-gw": {"client_id": "api-gw", "redirect_uris": ["http://localhost:3000/callback", "https://api.riotpiao.com/callback"]},
    "minio": {"client_id": "minio", "redirect_uris": ["http://localhost:9000/auth/sso/oauth2/code", "https://minio.riotpiao.com/auth/sso/oauth2/code"]},
    "poimen": {"client_id": "poimen", "redirect_uris": ["http://localhost:3000/callback", "https://poimen.riotpiao.com/callback"]},
    "paperless": {"client_id": "paperless", "redirect_uris": ["http://localhost:8000/auth/complete", "https://paperless.riotpiao.com/auth/complete"]},
    "grafana": {"client_id": "grafana", "redirect_uris": ["http://localhost:3000/login/generic_oauth", "https://grafana.riotpiao.com/login/generic_oauth"]},
}

for provider_name, provider_spec in OAuth2_PROVIDERS.items():
    if provider_name in existing_providers:
        print(f"  {provider_name}: already exists")
    else:
        # Build redirect_uris list with proper schema
        redirect_uris_list = [{"url": uri, "matching_mode": "strict"} for uri in provider_spec["redirect_uris"]]
        
        client_secret = os.environ.get(f"AUTHENTIK_PROVIDER_{provider_name.upper()}_SECRET", f"{provider_name}-secret-placeholder")
        
        status, res = api("POST", "/api/v3/providers/oauth2/", {
            "name": provider_name,
            "authorization_flow": auth_flow,
            "invalidation_flow": inval_flow,
            "grant_types": ["authorization_code", "implicit", "password"],
            "client_id": provider_spec["client_id"],
            "client_secret": client_secret,
            "redirect_uris": redirect_uris_list,
            "property_mappings": scope_pks_list,
        })
        if status in (200, 201):
            print(f"  {provider_name}: created")
        else:
            print(f"  {provider_name}: FAILED {status} {res}")

# ===========================================================================
# Phase 6: Create OAuth2 Applications (bind providers to public token endpoints)
# ===========================================================================
print("\n[6/7] Creating OAuth2 Applications...")
print("     (binds providers to /application/o/token/ endpoints)")

status, res = api("GET", "/api/v3/core/applications/?page_size=100")
if status != 200:
    die(f"GET applications -> {status} {res}")
existing_apps = {a["slug"]: a for a in res.get("results", [])}

for provider_name in OAuth2_PROVIDERS.keys():
    if provider_name in existing_apps:
        print(f"  {provider_name}: already exists")
    else:
        # Get the provider PK to link
        status, provider_res = api("GET", f"/api/v3/providers/oauth2/?name={provider_name}")
        if status != 200 or not provider_res.get("results"):
            print(f"  {provider_name}: provider not found")
            continue
        
        provider_pk = provider_res["results"][0]["pk"]
        
        status, res = api("POST", "/api/v3/core/applications/", {
            "name": provider_name,
            "slug": provider_name,
            "provider": provider_pk,
        })
        if status in (200, 201):
            print(f"  {provider_name}: created")
        else:
            print(f"  {provider_name}: FAILED {status} {res}")

# ===========================================================================
# Phase 7: Create/update rock user → homelab-admins, matching Forgejo identity
# ===========================================================================
print("\n[7/10] Creating/updating rock user (locartrock@gmail.com)...")

ROCK_EMAIL = "locartrock@gmail.com"
ROCK_PASSWORD = os.environ.get("ROCK_PASSWORD", "")

status, res = api("GET", "/api/v3/core/users/?username=rock")
if status == 200 and res.get("results"):
    rock_user = res["results"][0]
    # Ensure email matches Forgejo's rock user for OIDC linking
    patch_data = {"email": ROCK_EMAIL, "name": "Rock"}
    status, res = api("PATCH", f"/api/v3/core/users/{rock_user['pk']}/", patch_data)
    if status in (200, 201):
        print(f"  rock: updated email to {ROCK_EMAIL}")
    else:
        print(f"  rock: update FAILED {status} {res}")
else:
    if not ROCK_PASSWORD:
        print("  rock: NOT FOUND and ROCK_PASSWORD not set, skipping creation")
        print("    export ROCK_PASSWORD=<password> and re-run")
        rock_user = None
    else:
        status, res = api("POST", "/api/v3/core/users/", {
            "username": "rock",
            "name": "Rock",
            "email": ROCK_EMAIL,
            "is_active": True,
            "is_superuser": False,
            "password": ROCK_PASSWORD,
        })
        if status in (200, 201):
            rock_user = res
            print(f"  rock: created with email {ROCK_EMAIL}")
        else:
            print(f"  rock: create FAILED {status} {res}")
            rock_user = None

if rock_user:
    status, res = api("GET", "/api/v3/core/groups/?name=homelab-admins")
    if status == 200 and res.get("results"):
        admins_group = res["results"][0]
        status, res = api("POST", f"/api/v3/core/groups/{admins_group['pk']}/users/add/", {"pk": rock_user["pk"]})
        if status in (200, 201, 204):
            print(f"  rock: added to homelab-admins")
        else:
            print(f"  rock: group add {status} {res}")

# ===========================================================================
# Phase 8: Email recovery flow (password reset via email)
# ===========================================================================
print("\n[8/10] Creating email recovery flow...")

# Read SMTP config from gotify-smtp secret (same Gmail creds)
SMTP_HOST = "smtp.gmail.com"
SMTP_PORT = 587
SMTP_USER = "locartrock@gmail.com"
SMTP_FROM = "authentik@riotpiao.com"
# Password read from env at runtime: AUTHENTIK_EMAIL__PASSWORD

# 8a. Create email stage for recovery
status, res = api("GET", "/api/v3/stages/email/?name=email-recovery")
if status == 200 and res.get("results"):
    email_stage_pk = res["results"][0]["pk"]
    print("  email-recovery stage: already exists")
else:
    status, res = api("POST", "/api/v3/stages/email/", {
        "name": "email-recovery",
        "use_global_settings": False,
        "host": SMTP_HOST,
        "port": SMTP_PORT,
        "username": SMTP_USER,
        "password": os.environ.get("AUTHENTIK_EMAIL_PASSWORD", ""),
        "use_tls": True,
        "use_ssl": False,
        "timeout": 10,
        "from_address": SMTP_FROM,
        "template": "email/password_reset.html",
        "activate_user_on_success": True,
    })
    if status in (200, 201):
        email_stage_pk = res["pk"]
        print("  email-recovery stage: created")
    else:
        email_stage_pk = None
        print(f"  email-recovery stage: FAILED {status} {res}")

# 8b. Create identification stage for recovery (email lookup)
status, res = api("GET", "/api/v3/stages/identification/?name=recovery-identification")
if status == 200 and res.get("results"):
    ident_stage_pk = res["results"][0]["pk"]
    print("  recovery-identification stage: already exists")
else:
    status, res = api("POST", "/api/v3/stages/identification/", {
        "name": "recovery-identification",
        "user_fields": ["email", "username"],
    })
    if status in (200, 201):
        ident_stage_pk = res["pk"]
        print("  recovery-identification stage: created")
    else:
        ident_stage_pk = None
        print(f"  recovery-identification stage: FAILED {status} {res}")

# 8c. Create password stage for new password entry
status, res = api("GET", "/api/v3/stages/password/?name=recovery-password-change")
if status == 200 and res.get("results"):
    pw_stage_pk = res["results"][0]["pk"]
    print("  recovery-password-change stage: already exists")
else:
    # Use prompt stage for password change instead
    status, res = api("GET", "/api/v3/stages/user_write/?name=recovery-user-write")
    if status == 200 and res.get("results"):
        pw_stage_pk = res["results"][0]["pk"]
        print("  recovery-user-write stage: already exists")
    else:
        status, res = api("POST", "/api/v3/stages/user_write/", {
            "name": "recovery-user-write",
        })
        if status in (200, 201):
            pw_stage_pk = res["pk"]
            print("  recovery-user-write stage: created")
        else:
            pw_stage_pk = None
            print(f"  recovery-user-write stage: FAILED {status} {res}")

# 8d. Create recovery flow
status, res = api("GET", "/api/v3/flows/instances/?slug=password-recovery")
if status == 200 and res.get("results"):
    recovery_flow_pk = res["results"][0]["pk"]
    print("  password-recovery flow: already exists")
else:
    status, res = api("POST", "/api/v3/flows/instances/", {
        "name": "Password Recovery",
        "slug": "password-recovery",
        "title": "Reset your password",
        "designation": "recovery",
    })
    if status in (200, 201):
        recovery_flow_pk = res["pk"]
        print("  password-recovery flow: created")
    else:
        recovery_flow_pk = None
        print(f"  password-recovery flow: FAILED {status} {res}")

# 8e. Bind stages to flow in order
if recovery_flow_pk and ident_stage_pk and email_stage_pk:
    for order, stage_pk, label in [
        (10, ident_stage_pk, "identification"),
        (20, email_stage_pk, "email"),
    ]:
        status, res = api("POST", "/api/v3/flows/bindings/", {
            "target": recovery_flow_pk,
            "stage": stage_pk,
            "order": order,
        })
        if status in (200, 201):
            print(f"  bound {label} stage at order {order}")
        elif status == 400 and "already exists" in str(res).lower():
            print(f"  {label} stage: already bound")
        else:
            print(f"  bind {label}: {status} {res}")

# ===========================================================================
# Phase 9: Set recovery flow on brand
# ===========================================================================
print("\n[9/10] Setting recovery flow on brand...")

if recovery_flow_pk:
    status, res = api("GET", "/api/v3/brands/instances/")
    if status == 200 and res.get("results"):
        brand = res["results"][0]
        status, res = api("PATCH", f"/api/v3/brands/instances/{brand['brand_uuid']}/", {
            "flow_recovery": recovery_flow_pk,
        })
        if status in (200, 201):
            print("  recovery flow set on brand")
        else:
            print(f"  FAILED {status} {res}")
    else:
        print("  no brand found")

# ===========================================================================
# Phase 10: Ensure Forgejo OAuth2 source uses matching email claim
# ===========================================================================
print("\n[10/10] Verifying Forgejo OIDC linkage...")
print(f"  rock@Authentik email: {ROCK_EMAIL}")
print("  Forgejo OIDC will match on email — ensure Forgejo's rock user")
print(f"  has email {ROCK_EMAIL} in Forgejo settings → Profile")

# ===========================================================================
# Summary
# ===========================================================================
print("\n" + "="*70)
print("AUTHENTIK PROVISIONING COMPLETE")
print("="*70)
print(f"\n  [1] Groups: {len(GROUPS)}")
print(f"  [2] Service accounts: {len(SERVICE_ACCOUNTS)}")
print(f"  [3] Scope mappings: {len(SCOPE_MAPPINGS)}")
print(f"  [4] Flows resolved")
print(f"  [5] OAuth2 providers: {len(OAuth2_PROVIDERS)}")
print(f"  [6] OAuth2 applications bound")
print(f"  [7] rock user (locartrock@gmail.com) -> homelab-admins")
print(f"  [8] Email recovery flow (smtp.gmail.com)")
print(f"  [9] Recovery flow set on brand")
print(f"  [10] Forgejo OIDC linkage verified")
print("\nNEXT: Set Forgejo rock user email to locartrock@gmail.com in Forgejo profile")
print("TEST: https://authentik.riotpiao.com/if/flow/password-recovery/")
print("="*70)