#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "azure-mgmt-cosmosdb",
#   "azure-identity",
# ]
# ///
"""
setup_mathblob.py — Configure Azure Cosmos DB for Gremlin (mathblob)
Reads role-assignments JSON to pull subscription/principal IDs,
then creates the database + graph and writes connection info.

Run with:  uv run scripts/setup_mathblob.py
"""

import json
import sys
from pathlib import Path
from azure.identity import InteractiveBrowserCredential
from azure.mgmt.cosmosdb import CosmosDBManagementClient
from azure.mgmt.cosmosdb.models import (
    GremlinDatabaseCreateUpdateParameters,
    GremlinDatabaseResource,
    GremlinGraphCreateUpdateParameters,
    GremlinGraphResource,
    ContainerPartitionKey,
    CreateUpdateOptions,  # used for graph (no throughput on serverless)
)

ROLE_FILE = Path.home() / "Downloads" / "role-assignments-2026-06-18.json"
ACCOUNT   = "mathblob"
RG        = "Mathblob"
DB_NAME   = "research"
GRAPH_NAME = "concepts"
PARTITION_KEY = "/pk"
# Serverless account — no throughput setting needed

# ── 1. Load role assignments ──────────────────────────────────────────────────

with open(ROLE_FILE) as f:
    roles = json.load(f)

me = roles[0]
subscription_id = me["Scope"].split("/subscriptions/")[1].split("/")[0]
principal_id    = me["ObjectId"]

print(f"Subscription : {subscription_id}")
print(f"Principal ID : {principal_id}")
print(f"Account      : {ACCOUNT}")
print()

# ── 2. Authenticate (opens browser once, caches token) ───────────────────────

print("Authenticating — browser will open if no cached token...")
credential = InteractiveBrowserCredential(tenant_id="00c36292-f660-4f91-8eed-1dfa0013060c")
client = CosmosDBManagementClient(credential, subscription_id)

# ── 3. Verify account is ready ────────────────────────────────────────────────

print("Checking Cosmos DB account status...")
account = client.database_accounts.get(RG, ACCOUNT)
state = account.provisioning_state
print(f"Provisioning state: {state}")
if state != "Succeeded":
    print("Account not ready yet — wait for deployment to complete then re-run.")
    sys.exit(0)

gremlin_endpoint = f"wss://{ACCOUNT}.gremlin.cosmos.azure.com:443/"
print(f"Gremlin endpoint: {gremlin_endpoint}")
print()

# ── 4. Create Gremlin database ────────────────────────────────────────────────

print(f"Creating Gremlin database '{DB_NAME}'...")
client.gremlin_resources.begin_create_update_gremlin_database(
    RG, ACCOUNT, DB_NAME,
    GremlinDatabaseCreateUpdateParameters(
        resource=GremlinDatabaseResource(id=DB_NAME),
        options=CreateUpdateOptions(),
    ),
).result()
print("  done.")

# ── 5. Create Gremlin graph ───────────────────────────────────────────────────

print(f"Creating Gremlin graph '{GRAPH_NAME}' (partition key: {PARTITION_KEY})...")
client.gremlin_resources.begin_create_update_gremlin_graph(
    RG, ACCOUNT, DB_NAME, GRAPH_NAME,
    GremlinGraphCreateUpdateParameters(
        resource=GremlinGraphResource(
            id=GRAPH_NAME,
            partition_key=ContainerPartitionKey(paths=[PARTITION_KEY]),
        ),
        options=CreateUpdateOptions(),
    ),
).result()
print("  done.")

# ── 6. Get primary key ────────────────────────────────────────────────────────

print("Fetching keys...")
keys = client.database_accounts.list_keys(RG, ACCOUNT)
primary_key = keys.primary_master_key
print("  done.")

# ── 6. Write .env file ────────────────────────────────────────────────────────

env_path = Path(__file__).parent.parent / ".env.gremlin"
env_content = f"""\
# mathblob Gremlin connection — DO NOT COMMIT
GREMLIN_ENDPOINT={gremlin_endpoint}
GREMLIN_USERNAME=/dbs/{DB_NAME}/colls/{GRAPH_NAME}
GREMLIN_PASSWORD={primary_key}
GREMLIN_DATABASE={DB_NAME}
GREMLIN_GRAPH={GRAPH_NAME}
AZURE_SUBSCRIPTION_ID={subscription_id}
AZURE_PRINCIPAL_ID={principal_id}
"""
env_path.write_text(env_content)
print(f"\nConnection info written to: {env_path}")

# ── 7. Print connection snippet ───────────────────────────────────────────────

print("""
── Python connection snippet ──────────────────────────────────────────────────
from gremlin_python.driver import client, serializer
import os

c = client.Client(
    os.environ["GREMLIN_ENDPOINT"],
    "g",
    username=os.environ["GREMLIN_USERNAME"],
    password=os.environ["GREMLIN_PASSWORD"],
    message_serializer=serializer.GraphSONSerializersV2d0(),
)

# smoke test
result = c.submit("g.V().count()").all().result()
print("Vertex count:", result)
──────────────────────────────────────────────────────────────────────────────
""")
print("Setup complete.")