#!/usr/bin/env python3
import os
import json
import urllib.parse
import urllib.request
import duckdb

ACCOUNT = os.environ.get("AZURE_STORAGE_ACCOUNT", "o11yprodjapaneast")
CLIENT_ID = os.environ.get("AZURE_CLIENT_ID", "")
TENANT_ID = os.environ.get("AZURE_TENANT_ID", "")
FED_TOKEN_FILE = os.environ.get("AZURE_FEDERATED_TOKEN_FILE", "")

# abfss path (same data as az://)
CONTAINER = "o11y-prod-shared-eastus2"
REL_PATH = "deltalake/1372813089189741281/7553f5ff-4163-4424-a0b4-ca774600a4d9/slowlogs"
table_path = f"abfss://{CONTAINER}@{ACCOUNT}.dfs.core.windows.net/{REL_PATH}"

if not ACCOUNT:
    raise SystemExit("Please set AZURE_STORAGE_ACCOUNT")
if not CLIENT_ID or not TENANT_ID or not FED_TOKEN_FILE:
    raise SystemExit("Please set AZURE_CLIENT_ID, AZURE_TENANT_ID, AZURE_FEDERATED_TOKEN_FILE")


def get_access_token():
    with open(FED_TOKEN_FILE, "r", encoding="utf-8") as f:
        federated_token = f.read().strip()
    token_url = f"https://login.microsoftonline.com/{TENANT_ID}/oauth2/v2.0/token"
    data = urllib.parse.urlencode(
        {
            "client_id": CLIENT_ID,
            "scope": "https://storage.azure.com/.default",
            "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
            "client_assertion": federated_token,
            "grant_type": "client_credentials",
        }
    ).encode("utf-8")
    req = urllib.request.Request(
        token_url,
        data=data,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        body = resp.read()
        if resp.status != 200:
            raise RuntimeError(f"Token request failed {resp.status}: {body.decode('utf-8')}")
    token_resp = json.loads(body.decode("utf-8"))
    access_token = token_resp.get("access_token")
    if not access_token:
        raise RuntimeError("access_token missing in response")
    return access_token



con = duckdb.connect()

# Load Azure/delta extensions
con.execute("INSTALL azure;")
con.execute("LOAD azure;")
con.execute("INSTALL delta;")
con.execute("LOAD delta;")

access_token = get_access_token()
print("access_token:", access_token)
con.execute(
    f"""
    CREATE OR REPLACE SECRET azure_default (
        TYPE AZURE,
        PROVIDER access_token,
        ACCOUNT_NAME '{ACCOUNT}',
        access_token '{access_token}'
    );
"""
)



con.execute("SET threads=5;")
con.execute("SET azure_transport_option_type=curl;")
con.execute("SET azure_read_transfer_concurrency=20;")

# Attach table using attach mode (same as Go code)
table_alias = "slowlogs_attached"
print(f"Attaching table: {table_path} as {table_alias}")
try:
    con.execute(f"attach '{table_path}' as {table_alias}(type delta);")
    print(f"Successfully attached table as {table_alias}")
except Exception as e:
    if "already exists" in str(e).lower() or "already attached" in str(e).lower():
        print(f"Table {table_alias} already attached, continuing...")
    else:
        raise

# Query using attached table (same pattern as Go code)
query = f"""
SELECT id, conn_id, digest, time, query_time, mem_max, query, instance, succ, (COALESCE(request_unit_read, 0) + COALESCE(request_unit_write, 0)) AS ru, db, is_internal 
FROM {table_alias} 
WHERE time >= 1765867996.000000 AND time <= 1765869797.000000 
  AND date >= '20251216' AND date <= '20251216' 
  AND (is_internal = false OR is_internal IS NULL) 
ORDER BY time DESC 
LIMIT 100;
"""

print(f"\nQuery using attached table {table_alias}:")
print("Query:", query)
df = con.execute(query).fetchdf()
print(df)