# Databricks notebook source
# MAGIC %md
# MAGIC # Sync Discovered Agents: Delta → Lakebase
# MAGIC
# MAGIC Reads the latest snapshot from the Delta `discovered_agents` table
# MAGIC and upserts into the Lakebase PostgreSQL `discovered_agents` table.
# MAGIC
# MAGIC This runs as the second task in the discovery Workflow, after
# MAGIC `01_discover_agents` has written the Delta table.

# COMMAND ----------

# MAGIC %pip install psycopg2-binary databricks-sdk requests --upgrade
# MAGIC dbutils.library.restartPython()

# COMMAND ----------

import json
import os
import psycopg2
from psycopg2.extras import execute_values
from datetime import datetime, timezone

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Configuration

# COMMAND ----------

dbutils.widgets.text("catalog", "", "Unity Catalog name")
dbutils.widgets.text("schema", "", "Schema name")
dbutils.widgets.text("delta_table", "discovered_agents", "Delta table name")
dbutils.widgets.text("lakebase_dns", "", "Lakebase host (DNS)")
dbutils.widgets.text("lakebase_database", "", "Lakebase database name")
dbutils.widgets.text("lakebase_instance", "", "Lakebase instance name (Provisioned only)")
dbutils.widgets.text("lakebase_endpoint_path", "", "Lakebase endpoint path (Autoscaling only, e.g. projects/<name>/branches/<branch>/endpoints/<endpoint>)")
dbutils.widgets.text("account_id", "", "Databricks account ID")
dbutils.widgets.text("warehouse_id", "", "SQL warehouse ID for system table queries")

CATALOG = dbutils.widgets.get("catalog")
SCHEMA = dbutils.widgets.get("schema")
DELTA_TABLE = f"{CATALOG}.{SCHEMA}.{dbutils.widgets.get('delta_table')}"
LAKEBASE_DNS = dbutils.widgets.get("lakebase_dns")
LAKEBASE_DATABASE = dbutils.widgets.get("lakebase_database")
LAKEBASE_INSTANCE = dbutils.widgets.get("lakebase_instance")
LAKEBASE_ENDPOINT_PATH = dbutils.widgets.get("lakebase_endpoint_path")
ACCOUNT_ID = dbutils.widgets.get("account_id")
WAREHOUSE_ID = dbutils.widgets.get("warehouse_id")

# Set as env var so refresh_workspace_registry can find it
if ACCOUNT_ID:
    os.environ["DATABRICKS_ACCOUNT_ID"] = ACCOUNT_ID

if not CATALOG or not SCHEMA:
    raise ValueError(
        f"catalog and schema must be set via job parameters (got catalog={CATALOG!r}, schema={SCHEMA!r}). "
        "Deploy with: databricks bundle deploy -t <target>"
    )
if not LAKEBASE_DNS:
    raise ValueError(
        f"lakebase_dns must be set via job parameters (got dns={LAKEBASE_DNS!r}). "
        "Deploy with: databricks bundle deploy -t <target>"
    )
if not LAKEBASE_ENDPOINT_PATH and not LAKEBASE_INSTANCE:
    raise ValueError(
        "Either lakebase_endpoint_path (Autoscaling) or lakebase_instance (Provisioned) must be set via job parameters."
    )

print(f"Source Delta table: {DELTA_TABLE}")
print(f"Target Lakebase: {LAKEBASE_DNS}/{LAKEBASE_DATABASE}")

# COMMAND ----------

# MAGIC %md
# MAGIC ## Generate Lakebase Credentials

# COMMAND ----------

def get_lakebase_connection():
    """Create a Lakebase PostgreSQL connection.

    Tries the Databricks SDK first (w.database.generate_database_credential),
    then falls back to the REST API if the SDK version is too old.
    """
    from databricks.sdk import WorkspaceClient
    import requests as http_requests
    import uuid

    w = WorkspaceClient()

    # Get current identity for PG username
    me = w.current_user.me()
    pg_user = me.user_name
    print(f"Lakebase user: {pg_user}")

    # Get auth token for REST API calls
    def _get_auth_token():
        try:
            header_factory = w.config.authenticate
            auth_headers = header_factory()
            return auth_headers.get("Authorization", "").replace("Bearer ", "")
        except Exception:
            pass
        try:
            return dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()
        except Exception:
            return getattr(w.config, "token", "") or ""

    pg_password = None
    host = w.config.host.rstrip("/")

    # Autoscaling Lakebase: POST /api/2.0/postgres/credentials with explicit endpoint path
    if LAKEBASE_ENDPOINT_PATH:
        try:
            token = _get_auth_token()
            headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
            cred_resp = http_requests.post(
                f"{host}/api/2.0/postgres/credentials",
                headers=headers,
                json={"endpoint": LAKEBASE_ENDPOINT_PATH},
            )
            cred_resp.raise_for_status()
            pg_password = cred_resp.json().get("token", "")
            if pg_password:
                print(f"Credential generated via Autoscaling API (endpoint={LAKEBASE_ENDPOINT_PATH})")
        except Exception as e:
            print(f"Autoscaling credential generation failed: {e}")

    # Provisioned Lakebase SDK (w.database.generate_database_credential)
    if not pg_password and LAKEBASE_INSTANCE and hasattr(w, "database"):
        try:
            creds = w.database.generate_database_credential(
                instance_names=[LAKEBASE_INSTANCE]
            )
            pg_password = creds.token
            print("Credential generated via Provisioned SDK")
        except Exception as e:
            print(f"Provisioned SDK credential generation failed: {e}")

    # Fallback to Provisioned REST API
    if not pg_password and LAKEBASE_INSTANCE:
        print("Falling back to Provisioned REST API for credential generation")
        token = _get_auth_token()
        headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
        cred_resp = http_requests.post(
            f"{host}/api/2.0/database/credentials",
            headers=headers,
            json={"instance_names": [LAKEBASE_INSTANCE], "request_id": str(uuid.uuid4())},
        )
        cred_resp.raise_for_status()
        pg_password = cred_resp.json().get("token", "")
        print("Credential generated via Provisioned REST API")

    if not pg_password:
        raise RuntimeError("Could not generate Lakebase credentials via any method")

    print(f"Connecting to Lakebase at: {LAKEBASE_DNS}")

    return psycopg2.connect(
        host=LAKEBASE_DNS,
        port=5432,
        database=LAKEBASE_DATABASE,
        user=pg_user,
        password=pg_password,
        sslmode="require",
        connect_timeout=15,
        options="-c statement_timeout=60000",
    )

# COMMAND ----------

# MAGIC %md
# MAGIC ## Ensure Lakebase Table Exists

# COMMAND ----------

def ensure_lakebase_table(conn):
    """Create the discovered_agents table in Lakebase if it doesn't exist."""
    ddl = """
    CREATE TABLE IF NOT EXISTS discovered_agents (
        agent_id          TEXT PRIMARY KEY,
        workspace_id      TEXT NOT NULL,
        name              TEXT NOT NULL,
        type              TEXT,
        endpoint_name     TEXT,
        endpoint_status   TEXT,
        model_name        TEXT,
        served_entity_name TEXT,
        creator           TEXT,
        description       TEXT DEFAULT '',
        config            JSONB,
        last_synced       TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
        source            TEXT DEFAULT 'api',
        is_extensive      BOOLEAN DEFAULT FALSE
    );
    CREATE INDEX IF NOT EXISTS idx_da_ws ON discovered_agents (workspace_id);
    CREATE INDEX IF NOT EXISTS idx_da_type ON discovered_agents (type);
    DROP INDEX IF EXISTS idx_da_name_ws;
    CREATE INDEX IF NOT EXISTS idx_da_name_ws ON discovered_agents (name, workspace_id);
    """
    with conn.cursor() as cur:
        for stmt in ddl.split(";"):
            stmt = stmt.strip()
            if stmt:
                try:
                    cur.execute(stmt)
                except Exception as e:
                    print(f"  DDL warning: {e}")
    conn.commit()
    print("Lakebase table ensured")

# COMMAND ----------

# MAGIC %md
# MAGIC ## Read Delta & Sync

# COMMAND ----------

# Read the latest snapshot from Delta
df = spark.read.table(DELTA_TABLE)
agent_count = df.count()
print(f"Read {agent_count} agents from Delta table")

if agent_count == 0:
    print("No agents to sync — exiting")
    dbutils.notebook.exit(json.dumps({"status": "skipped", "reason": "no_agents", "count": 0}))

# Collect to driver (discovery tables are small — typically < 1000 rows)
rows = df.collect()

# COMMAND ----------

# Connect to Lakebase and sync
conn = get_lakebase_connection()

try:
    ensure_lakebase_table(conn)

    with conn.cursor() as cur:
        # Drop the UNIQUE constraint on (name, workspace_id) that causes conflicts
        # with duplicate names from different discovery sources (e.g. audit log Genie Spaces).
        # Try multiple approaches since it could be an index or a table constraint.
        for drop_sql in [
            "DROP INDEX IF EXISTS idx_da_name_ws",
            "ALTER TABLE discovered_agents DROP CONSTRAINT IF EXISTS idx_da_name_ws",
        ]:
            try:
                cur.execute(drop_sql)
                conn.commit()
                print(f"  Executed: {drop_sql}")
            except Exception as e:
                conn.rollback()
                print(f"  {drop_sql} — skipped: {e}")

        # List remaining indexes/constraints for debugging
        try:
            cur.execute("""
                SELECT indexname, indexdef
                FROM pg_indexes
                WHERE tablename = 'discovered_agents'
            """)
            for row in cur.fetchall():
                print(f"  Index: {row[0]} -> {row[1]}")
        except Exception as e:
            print(f"  Could not list indexes: {e}")

        # Truncate and reload for a clean snapshot (matches the app's refresh_discovery behavior)
        cur.execute("TRUNCATE TABLE discovered_agents")
        conn.commit()
        print("Truncated existing Lakebase discovered_agents")

        # Bulk insert using execute_values with ON CONFLICT to handle
        # duplicate (name, workspace_id) pairs from different discovery sources
        insert_sql = """
            INSERT INTO discovered_agents
                (agent_id, workspace_id, name, type, endpoint_name,
                 endpoint_status, model_name, served_entity_name,
                 creator, description, config, last_synced, source, is_extensive)
            VALUES %s
            ON CONFLICT (agent_id) DO UPDATE SET
                name = EXCLUDED.name,
                type = EXCLUDED.type,
                endpoint_name = EXCLUDED.endpoint_name,
                endpoint_status = EXCLUDED.endpoint_status,
                model_name = EXCLUDED.model_name,
                served_entity_name = EXCLUDED.served_entity_name,
                creator = EXCLUDED.creator,
                description = EXCLUDED.description,
                config = EXCLUDED.config,
                last_synced = EXCLUDED.last_synced,
                source = EXCLUDED.source,
                is_extensive = EXCLUDED.is_extensive
        """

        values = []
        now = datetime.now(timezone.utc)
        for r in rows:
            config_val = r.config
            # Ensure config is valid JSON for JSONB column
            if config_val:
                try:
                    json.loads(config_val)
                except (json.JSONDecodeError, TypeError):
                    config_val = json.dumps({})
            else:
                config_val = json.dumps({})

            values.append((
                r.agent_id,
                r.workspace_id,
                r.name,
                r.type,
                r.endpoint_name or "",
                r.endpoint_status or "",
                r.model_name or "",
                r.served_entity_name or "",
                r.creator or "",
                r.description or "",
                config_val,
                now,
                r.source or "api",
                bool(r.is_extensive),
            ))

        execute_values(cur, insert_sql, values, page_size=100)
        conn.commit()

        # Verify
        cur.execute("SELECT COUNT(*) FROM discovered_agents")
        lb_count = cur.fetchone()[0]
        print(f"Synced {lb_count} agents to Lakebase")

finally:
    conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Sync Summary

# COMMAND ----------

agent_result = {
    "agents_delta_count": agent_count,
    "agents_lakebase_count": len(values),
}
print(json.dumps(agent_result, indent=2))

# COMMAND ----------

# MAGIC %md
# MAGIC # Phase 2: Sync MLflow Observability → Lakebase
# MAGIC
# MAGIC Three data sources:
# MAGIC - **System tables** (`system.mlflow.*`): experiments & runs (account-level, cross-workspace)
# MAGIC - **Delta table** (`observability_traces`): traces discovered by `04_discover_observability` task
# MAGIC
# MAGIC Each row is tagged with `data_source` for attribution.

# COMMAND ----------

import time as _time

# COMMAND ----------

# MAGIC %md
# MAGIC ## SQL Statements API helper (for system tables)

# COMMAND ----------

def _find_warehouse_id():
    """Find a running SQL warehouse. Uses job parameter first, then SDK lookup."""
    if WAREHOUSE_ID:
        print(f"  Using warehouse from parameter: {WAREHOUSE_ID}")
        return WAREHOUSE_ID
    from databricks.sdk import WorkspaceClient as _WRC3
    w = _WRC3()
    try:
        warehouses = list(w.warehouses.list())
        for wh in warehouses:
            if wh.warehouse_type and "SERVERLESS" in str(wh.warehouse_type).upper() and str(wh.state) == "RUNNING":
                print(f"  Found serverless warehouse: {wh.id}")
                return wh.id
        for wh in warehouses:
            if str(wh.state) == "RUNNING":
                print(f"  Found warehouse: {wh.id}")
                return wh.id
    except Exception as exc:
        print(f"  ⚠️  Warehouse lookup failed: {exc}")
    return None


def _execute_system_sql(sql):
    """Execute SQL via SQL Statements API and return list of dicts."""
    wh_id = _find_warehouse_id()
    if not wh_id:
        print("  ⚠️  No running SQL warehouse found — cannot query system tables")
        return []

    from databricks.sdk import WorkspaceClient as _WRC4
    w = _WRC4()

    body = {
        "warehouse_id": wh_id,
        "statement": sql,
        "wait_timeout": "50s",
        "disposition": "INLINE",
        "format": "JSON_ARRAY",
    }

    try:
        resp_json = w.api_client.do("POST", "/api/2.0/sql/statements", body=body)
    except Exception as exc:
        print(f"  ⚠️  SQL Statements API failed: {exc}")
        return []

    status = resp_json.get("status", {}).get("state", "")
    statement_id = resp_json.get("statement_id", "")

    # Poll if still running
    if status in ("PENDING", "RUNNING") and statement_id:
        for _ in range(20):
            _time.sleep(3)
            try:
                resp_json = w.api_client.do("GET", f"/api/2.0/sql/statements/{statement_id}")
            except Exception:
                continue
            status = resp_json.get("status", {}).get("state", "")
            if status not in ("PENDING", "RUNNING"):
                break

    if status != "SUCCEEDED":
        err = resp_json.get("status", {}).get("error", {})
        print(f"  ⚠️  SQL status: {status} — {err.get('message', '')}")
        return []

    columns = [c["name"] for c in resp_json.get("manifest", {}).get("schema", {}).get("columns", [])]
    rows = resp_json.get("result", {}).get("data_array", [])
    return [dict(zip(columns, row)) for row in rows]

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 2a: System Tables — Experiments & Runs

# COMMAND ----------

print("▸ Querying system.mlflow.experiments_latest via SQL Statements API ...")
try:
    st_experiments = _execute_system_sql("""
        SELECT
            CAST(experiment_id AS STRING) AS experiment_id,
            CAST(workspace_id AS STRING) AS workspace_id,
            name,
            CASE WHEN delete_time IS NULL THEN 'active' ELSE 'deleted' END AS lifecycle_stage,
            CAST(UNIX_TIMESTAMP(COALESCE(update_time, create_time)) * 1000 AS BIGINT) AS last_update_time
        FROM system.mlflow.experiments_latest
        WHERE delete_time IS NULL
        ORDER BY COALESCE(update_time, create_time) DESC
        LIMIT 5000
    """)
    for e in st_experiments:
        e["data_source"] = "system_table"
    print(f"  ✅ {len(st_experiments)} experiments from system tables")
except Exception as exc:
    print(f"  ⚠️  system.mlflow.experiments_latest query failed: {exc}")
    st_experiments = []

# COMMAND ----------

print("▸ Querying system.mlflow.runs_latest via SQL Statements API ...")
try:
    st_runs = _execute_system_sql("""
        SELECT
            run_id,
            CAST(experiment_id AS STRING) AS experiment_id,
            CAST(workspace_id AS STRING) AS workspace_id,
            status,
            CAST(UNIX_TIMESTAMP(start_time) * 1000 AS BIGINT) AS start_time,
            CAST(UNIX_TIMESTAMP(end_time) * 1000 AS BIGINT) AS end_time,
            created_by AS user_id,
            run_name,
            TO_JSON(tags) AS tags,
            TO_JSON(params) AS params,
            TO_JSON(aggregated_metrics) AS metrics
        FROM system.mlflow.runs_latest
        WHERE delete_time IS NULL
        ORDER BY start_time DESC
        LIMIT 5000
    """)
    for r in st_runs:
        r["data_source"] = "system_table"
    print(f"  ✅ {len(st_runs)} runs from system tables")
except Exception as exc:
    print(f"  ⚠️  system.mlflow.runs_latest query failed: {exc}")
    st_runs = []

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 2b: Read Traces from Delta (populated by discover_observability task)

# COMMAND ----------

TRACES_DELTA = f"{CATALOG}.{SCHEMA}.observability_traces"
TRACE_DETAILS_DELTA = f"{CATALOG}.{SCHEMA}.observability_trace_details"
TRACES_UC_OTEL_DELTA = f"{CATALOG}.{SCHEMA}.observability_traces_uc_otel"
TRACE_DETAILS_UC_OTEL_DELTA = f"{CATALOG}.{SCHEMA}.observability_trace_details_uc_otel"
trace_count = 0
trace_detail_count = 0

print(f"▸ Reading traces from Delta: {TRACES_DELTA} ...")
try:
    traces_df = spark.read.table(TRACES_DELTA)
    delta_trace_count = traces_df.count()
    print(f"  ✅ {delta_trace_count} traces in Delta table (default-backend, REST)")
    delta_traces = traces_df.collect()
except Exception as exc:
    print(f"  ⚠️  Could not read traces Delta table: {exc}")
    delta_traces = []
    delta_trace_count = 0

print(f"▸ Reading UC-OTel traces from Delta: {TRACES_UC_OTEL_DELTA} ...")
try:
    uc_traces_df = spark.read.table(TRACES_UC_OTEL_DELTA)
    delta_uc_trace_count = uc_traces_df.count()
    print(f"  ✅ {delta_uc_trace_count} traces in Delta table (Tier 2b: UC OTel)")
    delta_uc_traces = uc_traces_df.collect()
except Exception as exc:
    print(f"  ⚠️  Could not read UC OTel traces Delta table: {exc}")
    delta_uc_traces = []
    delta_uc_trace_count = 0

print(f"▸ Reading trace details from Delta: {TRACE_DETAILS_DELTA} ...")
try:
    details_df = spark.read.table(TRACE_DETAILS_DELTA)
    delta_detail_count = details_df.count()
    print(f"  ✅ {delta_detail_count} trace details in Delta table (default-backend, REST)")
    delta_trace_details = details_df.collect()
except Exception as exc:
    print(f"  ⚠️  Could not read trace details Delta table: {exc}")
    delta_trace_details = []
    delta_detail_count = 0

print(f"▸ Reading UC-OTel trace details from Delta: {TRACE_DETAILS_UC_OTEL_DELTA} ...")
try:
    uc_details_df = spark.read.table(TRACE_DETAILS_UC_OTEL_DELTA)
    delta_uc_detail_count = uc_details_df.count()
    print(f"  ✅ {delta_uc_detail_count} trace details in Delta table (Tier 2b: UC OTel)")
    delta_uc_trace_details = uc_details_df.collect()
except Exception as exc:
    print(f"  ⚠️  Could not read UC OTel trace details Delta table: {exc}")
    delta_uc_trace_details = []
    delta_uc_detail_count = 0

# Combine REST and UC OTel rows for the upsert path below — both go to the
# same Lakebase tables, with `data_source`/`source_type` distinguishing them.
delta_traces = list(delta_traces) + list(delta_uc_traces)
delta_trace_details = list(delta_trace_details) + list(delta_uc_trace_details)
print(f"▸ Combined: {len(delta_traces)} trace rows, {len(delta_trace_details)} detail rows for Lakebase upsert")

# COMMAND ----------

# MAGIC %md
# MAGIC ## Upsert All Observability Data to Lakebase

# COMMAND ----------

obs_conn = get_lakebase_connection()
now = datetime.now(timezone.utc)

all_experiments = st_experiments

# Ensure observability tables
with obs_conn.cursor() as cur:
    for ddl in [
        """CREATE TABLE IF NOT EXISTS observability_experiments (
            experiment_id TEXT NOT NULL, workspace_id TEXT NOT NULL,
            name TEXT, lifecycle_stage TEXT, last_update_time BIGINT,
            artifact_location TEXT, tags JSONB,
            data_source TEXT DEFAULT 'system_table',
            last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (workspace_id, experiment_id))""",
        "CREATE INDEX IF NOT EXISTS idx_oe_ws ON observability_experiments (workspace_id)",
        "ALTER TABLE observability_experiments ADD COLUMN IF NOT EXISTS data_source TEXT DEFAULT 'system_table'",
        "ALTER TABLE observability_experiments ADD COLUMN IF NOT EXISTS tags JSONB",
        """CREATE TABLE IF NOT EXISTS observability_runs (
            run_id TEXT NOT NULL, workspace_id TEXT NOT NULL,
            experiment_id TEXT, status TEXT,
            start_time BIGINT, end_time BIGINT,
            user_id TEXT, run_name TEXT,
            tags JSONB, params JSONB, metrics JSONB,
            data_source TEXT DEFAULT 'system_table',
            last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (workspace_id, run_id))""",
        "CREATE INDEX IF NOT EXISTS idx_or_ws ON observability_runs (workspace_id)",
        "ALTER TABLE observability_runs ADD COLUMN IF NOT EXISTS data_source TEXT DEFAULT 'system_table'",
        "ALTER TABLE observability_runs ADD COLUMN IF NOT EXISTS tags JSONB",
        "ALTER TABLE observability_runs ADD COLUMN IF NOT EXISTS params JSONB",
        "ALTER TABLE observability_runs ADD COLUMN IF NOT EXISTS metrics JSONB",
        """CREATE TABLE IF NOT EXISTS observability_traces (
            request_id TEXT NOT NULL, workspace_id TEXT NOT NULL,
            experiment_id TEXT, trace_name TEXT, state TEXT,
            request_time TEXT, execution_duration BIGINT,
            user_message TEXT, response_preview TEXT, token_usage JSONB,
            model_id TEXT, session_id TEXT, trace_user TEXT, source TEXT,
            tags JSONB, data_source TEXT DEFAULT 'rest_api',
            last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (workspace_id, request_id))""",
        "CREATE INDEX IF NOT EXISTS idx_ot_ws ON observability_traces (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_ot_time ON observability_traces (request_time DESC)",
        "ALTER TABLE observability_traces ADD COLUMN IF NOT EXISTS data_source TEXT DEFAULT 'rest_api'",
        """CREATE TABLE IF NOT EXISTS observability_trace_details (
            workspace_id TEXT NOT NULL, request_id TEXT NOT NULL,
            experiment_id TEXT, trace_info JSONB, trace_data JSONB,
            request_raw TEXT, response_raw TEXT,
            size_bytes INTEGER, source_type TEXT,
            cached_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (workspace_id, request_id))""",
        "CREATE INDEX IF NOT EXISTS idx_otd_ws     ON observability_trace_details (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_otd_cached ON observability_trace_details (cached_at DESC)",
    ]:
        # Each statement in its own savepoint — protects against the case
        # where the workflow runs as a non-owner of pre-existing tables (the
        # app's startup helpers create these as the SP). A failure rolls back
        # to the savepoint instead of poisoning the surrounding transaction.
        try:
            cur.execute("SAVEPOINT obs_ddl_sp")
            cur.execute(ddl)
            cur.execute("RELEASE SAVEPOINT obs_ddl_sp")
        except Exception as e:
            try: cur.execute("ROLLBACK TO SAVEPOINT obs_ddl_sp")
            except Exception: pass
            print(f"  DDL warning: {e}")
    obs_conn.commit()
print("✅ Observability tables ensured")

# COMMAND ----------

# Upsert experiments (REST data_source wins on conflict — richer info with tags)
exp_count = 0
if all_experiments:
    exp_values = []
    for exp in all_experiments:
        # REST API returns tags as array of {key, value} objects
        # System table experiments don't have tags
        tags_raw = exp.get("tags", [])
        if isinstance(tags_raw, list):
            # Convert [{key: k, value: v}, ...] to {k: v} dict for JSONB
            tags_dict = {t.get("key", ""): t.get("value", "") for t in tags_raw if isinstance(t, dict)}
            tags_json = json.dumps(tags_dict) if tags_dict else None
        elif isinstance(tags_raw, dict):
            tags_json = json.dumps(tags_raw) if tags_raw else None
        else:
            tags_json = None

        exp_values.append((
            str(exp.get("experiment_id", "")),
            str(exp.get("workspace_id", "")),
            exp.get("name", ""),
            exp.get("lifecycle_stage", ""),
            int(exp.get("last_update_time") or 0),
            exp.get("artifact_location", ""),
            tags_json,
            exp.get("data_source", "system_table"),
            now,
        ))
    with obs_conn.cursor() as cur:
        execute_values(
            cur,
            """INSERT INTO observability_experiments
               (experiment_id, workspace_id, name, lifecycle_stage, last_update_time, artifact_location, tags, data_source, last_synced)
               VALUES %s
               ON CONFLICT (workspace_id, experiment_id) DO UPDATE SET
                   name = EXCLUDED.name, lifecycle_stage = EXCLUDED.lifecycle_stage,
                   last_update_time = EXCLUDED.last_update_time,
                   artifact_location = EXCLUDED.artifact_location,
                   tags = COALESCE(EXCLUDED.tags, observability_experiments.tags),
                   data_source = EXCLUDED.data_source, last_synced = EXCLUDED.last_synced""",
            exp_values, page_size=100,
        )
        obs_conn.commit()
        exp_count = len(exp_values)
print(f"✅ Upserted {exp_count} experiments (system_table)")

# COMMAND ----------

# Upsert runs (system table — includes tags, params, metrics)
run_count = 0
if st_runs:
    run_values = []
    for r in st_runs:
        # tags/params/metrics come as JSON strings from SQL or dicts — normalize to JSON strings
        tags_val = r.get("tags", {})
        params_val = r.get("params", {})
        metrics_val = r.get("metrics", [])
        if isinstance(tags_val, str):
            tags_json = tags_val
        else:
            tags_json = json.dumps(tags_val) if tags_val else json.dumps({})
        if isinstance(params_val, str):
            params_json = params_val
        else:
            params_json = json.dumps(params_val) if params_val else json.dumps({})
        if isinstance(metrics_val, str):
            metrics_json = metrics_val
        else:
            metrics_json = json.dumps(metrics_val) if metrics_val else json.dumps([])

        run_values.append((
            r.get("run_id", ""),
            str(r.get("workspace_id", "")),
            str(r.get("experiment_id", "")),
            r.get("status", ""),
            int(r.get("start_time") or 0),
            int(r.get("end_time") or 0) if r.get("end_time") else None,
            r.get("user_id", ""),
            r.get("run_name", ""),
            tags_json,
            params_json,
            metrics_json,
            r.get("data_source", "system_table"),
            now,
        ))
    with obs_conn.cursor() as cur:
        execute_values(
            cur,
            """INSERT INTO observability_runs
               (run_id, workspace_id, experiment_id, status, start_time, end_time, user_id, run_name, tags, params, metrics, data_source, last_synced)
               VALUES %s
               ON CONFLICT (workspace_id, run_id) DO UPDATE SET
                   experiment_id = EXCLUDED.experiment_id, status = EXCLUDED.status,
                   start_time = EXCLUDED.start_time, end_time = EXCLUDED.end_time,
                   user_id = EXCLUDED.user_id, run_name = EXCLUDED.run_name,
                   tags = EXCLUDED.tags, params = EXCLUDED.params, metrics = EXCLUDED.metrics,
                   data_source = EXCLUDED.data_source, last_synced = EXCLUDED.last_synced""",
            run_values, page_size=100,
        )
        obs_conn.commit()
        run_count = len(run_values)
print(f"✅ Upserted {run_count} runs (system_table)")

# COMMAND ----------

# Upsert traces (from Delta table populated by discover_observability task)
# We enrich each slim trace row with fields parsed from the matching detail
# row's trace_info JSON — token_usage, user_message, response_preview — so the
# list endpoint surfaces them without an extra JSONB read per row.
trace_count = 0

# Build a lookup of detail rows by request_id (UC traces have empty workspace_id,
# default-backend traces are scoped — request_id is unique enough either way).
details_by_rid: dict = {}
for d in delta_trace_details:
    if d.request_id:
        details_by_rid[d.request_id] = d


def _enrich_from_trace_info(rid: str) -> tuple:
    """Return (token_usage_json, user_message, response_preview) for a trace."""
    d = details_by_rid.get(rid)
    if not d:
        return ("{}", "", "")
    ti_raw = d.trace_info or "{}"
    try:
        ti = json.loads(ti_raw) if isinstance(ti_raw, str) else (ti_raw or {})
    except Exception:
        ti = {}
    meta = ti.get("trace_metadata") or {}
    # token_usage lives under trace_metadata as a JSON-string value
    tu_str = meta.get("mlflow.trace.tokenUsage") or "{}"
    try:
        tu_obj = json.loads(tu_str) if isinstance(tu_str, str) else (tu_str or {})
    except Exception:
        tu_obj = {}
    tu_json = json.dumps(tu_obj) if isinstance(tu_obj, dict) else "{}"
    # Best-effort user message + response preview — populated by the parser for
    # default-backend traces (request_preview tag) and explicitly for UC traces.
    user_msg = ti.get("request_preview") or ""
    if not user_msg:
        # Fallback: pull first user-role text from the parsed request payload
        req_raw = d.request_raw or ""
        try:
            req_parsed = json.loads(req_raw) if req_raw else {}
            for inp in (req_parsed.get("input") or []):
                if isinstance(inp, dict) and inp.get("role") == "user":
                    c = inp.get("content")
                    if isinstance(c, str):
                        user_msg = c; break
                    if isinstance(c, list):
                        for piece in c:
                            if isinstance(piece, dict) and piece.get("text"):
                                user_msg = piece["text"]; break
                        if user_msg: break
        except Exception:
            pass
    resp_preview = ti.get("response_preview") or ""
    return (tu_json, user_msg or "", resp_preview or "")


if delta_traces:
    trace_values = []
    for r in delta_traces:
        rid = r.request_id
        if not rid:
            continue
        token_usage_json, user_msg, resp_preview = _enrich_from_trace_info(rid)
        trace_values.append((
            rid,
            r.workspace_id or "",
            r.experiment_id or "",
            r.trace_name or "",
            r.state or "",
            r.request_time or "",
            r.execution_duration,
            user_msg,
            resp_preview,
            token_usage_json,
            r.model_id or "",
            r.session_id or "",
            r.trace_user or "",
            r.source or "",
            r.tags or "{}",
            r.data_source or "rest_api",
            now,
        ))
    with obs_conn.cursor() as cur:
        execute_values(
            cur,
            """INSERT INTO observability_traces
               (request_id, workspace_id, experiment_id, trace_name, state,
                request_time, execution_duration, user_message, response_preview,
                token_usage, model_id, session_id, trace_user, source, tags,
                data_source, last_synced)
               VALUES %s
               ON CONFLICT (workspace_id, request_id) DO UPDATE SET
                   trace_name = EXCLUDED.trace_name, state = EXCLUDED.state,
                   request_time = EXCLUDED.request_time, execution_duration = EXCLUDED.execution_duration,
                   user_message = EXCLUDED.user_message, response_preview = EXCLUDED.response_preview,
                   token_usage = EXCLUDED.token_usage, model_id = EXCLUDED.model_id,
                   session_id = EXCLUDED.session_id, trace_user = EXCLUDED.trace_user,
                   source = EXCLUDED.source, tags = EXCLUDED.tags,
                   data_source = EXCLUDED.data_source, last_synced = EXCLUDED.last_synced""",
            trace_values, page_size=100,
        )
        obs_conn.commit()
        trace_count = len(trace_values)
print(f"✅ Upserted {trace_count} traces (from Delta, enriched from trace_info)")

# COMMAND ----------

# Upsert trace details (full spans + payloads) for cross-workspace cache.
# Note: UC-OTel traces have empty workspace_id (UC data is account-level, not
# workspace-scoped); we accept empty workspace_id here since the URI-form
# request_id is globally unique on its own.
if delta_trace_details:
    detail_values = []
    for r in delta_trace_details:
        rid = r.request_id
        if not rid:
            continue
        ws_id = r.workspace_id or ""
        detail_values.append((
            ws_id,
            rid,
            r.experiment_id or "",
            r.trace_info or "{}",
            r.trace_data or "{}",
            r.request_raw or "",
            r.response_raw or "",
            int(r.size_bytes or 0),
            r.source_type or "mlflow_rest",
            now,
        ))
    if detail_values:
        with obs_conn.cursor() as cur:
            execute_values(
                cur,
                """INSERT INTO observability_trace_details
                   (workspace_id, request_id, experiment_id, trace_info, trace_data,
                    request_raw, response_raw, size_bytes, source_type, cached_at)
                   VALUES %s
                   ON CONFLICT (workspace_id, request_id) DO UPDATE SET
                       experiment_id = EXCLUDED.experiment_id,
                       trace_info = EXCLUDED.trace_info,
                       trace_data = EXCLUDED.trace_data,
                       request_raw = EXCLUDED.request_raw,
                       response_raw = EXCLUDED.response_raw,
                       size_bytes = EXCLUDED.size_bytes,
                       source_type = EXCLUDED.source_type,
                       cached_at = EXCLUDED.cached_at""",
                detail_values, page_size=50,
            )
            obs_conn.commit()
            trace_detail_count = len(detail_values)
print(f"✅ Upserted {trace_detail_count} trace details (from Delta)")

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 2c: Sync AI Gateway / Inference Logs (Tier 2a, Delta → Lakebase)

# COMMAND ----------

GATEWAY_LOGS_DELTA = f"{CATALOG}.{SCHEMA}.gateway_inference_logs"
gw_log_count = 0

print(f"▸ Reading inference logs from Delta: {GATEWAY_LOGS_DELTA} ...")
try:
    gw_logs_df = spark.read.table(GATEWAY_LOGS_DELTA)
    gw_log_delta_count = gw_logs_df.count()
    print(f"  ✅ {gw_log_delta_count} inference-log rows in Delta")
    gw_log_rows = gw_logs_df.collect()
except Exception as exc:
    print(f"  ⚠️  Could not read inference logs Delta: {exc}")
    gw_log_rows = []
    gw_log_delta_count = 0

# Ensure Lakebase target table
with obs_conn.cursor() as cur:
    cur.execute("""
        CREATE TABLE IF NOT EXISTS gateway_inference_logs (
            request_id          TEXT NOT NULL,
            source_table        TEXT NOT NULL,
            client_request_id   TEXT,
            request_time        TIMESTAMP WITH TIME ZONE,
            status_code         INTEGER,
            execution_ms        BIGINT,
            request_payload     TEXT,
            response_payload    TEXT,
            request_size_bytes  BIGINT,
            response_size_bytes BIGINT,
            served_entity_id    TEXT,
            requester           TEXT,
            last_synced         TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (source_table, request_id)
        )
    """)
    # Extracted-payload columns — populated at sync time by parsing the JSON
    # payload bodies. Best-effort across vendor shapes (OpenAI/Anthropic/etc.).
    # Each ALTER runs in its own savepoint so a single column failure
    # doesn't poison the surrounding transaction (psycopg2 does not auto-
    # rollback on caught exceptions, and any subsequent statement on the
    # same cursor would otherwise hit InFailedSqlTransaction).
    for col_ddl in (
        "ALTER TABLE gateway_inference_logs ADD COLUMN IF NOT EXISTS model TEXT",
        "ALTER TABLE gateway_inference_logs ADD COLUMN IF NOT EXISTS input_tokens BIGINT",
        "ALTER TABLE gateway_inference_logs ADD COLUMN IF NOT EXISTS output_tokens BIGINT",
        "ALTER TABLE gateway_inference_logs ADD COLUMN IF NOT EXISTS total_tokens BIGINT",
        "ALTER TABLE gateway_inference_logs ADD COLUMN IF NOT EXISTS finish_reason TEXT",
        "ALTER TABLE gateway_inference_logs ADD COLUMN IF NOT EXISTS tool_call_count INTEGER",
    ):
        try:
            cur.execute("SAVEPOINT gw_col_sp")
            cur.execute(col_ddl)
            cur.execute("RELEASE SAVEPOINT gw_col_sp")
        except Exception as e:
            try: cur.execute("ROLLBACK TO SAVEPOINT gw_col_sp")
            except Exception: pass
            print(f"  gw col DDL warn: {e}")
    # CREATE INDEX IF NOT EXISTS still requires table ownership in PG, even
    # when the index already exists. Wrap each in a savepoint so a non-owner
    # workflow run doesn't fail when the indexes were pre-created (typically
    # by the app's startup ensure_*_table() helpers running as the SP).
    for idx_ddl in (
        "CREATE INDEX IF NOT EXISTS idx_gil_time   ON gateway_inference_logs (request_time DESC)",
        "CREATE INDEX IF NOT EXISTS idx_gil_source ON gateway_inference_logs (source_table)",
        "CREATE INDEX IF NOT EXISTS idx_gil_status ON gateway_inference_logs (status_code)",
        "CREATE INDEX IF NOT EXISTS idx_gil_model  ON gateway_inference_logs (model)",
    ):
        try:
            cur.execute("SAVEPOINT gw_idx_sp")
            cur.execute(idx_ddl)
            cur.execute("RELEASE SAVEPOINT gw_idx_sp")
        except Exception as e:
            try: cur.execute("ROLLBACK TO SAVEPOINT gw_idx_sp")
            except Exception: pass
            print(f"  gw idx DDL warn: {e}")
    obs_conn.commit()


def _parse_gateway_payload(req_raw: str, resp_raw: str) -> dict:
    """Best-effort extraction of model + token usage + finish reason.

    Handles OpenAI-style (`response.model`, `response.usage.{prompt,completion,total}_tokens`)
    and Anthropic-passthrough variants (`completion_tokens` ↔ `output_tokens`).
    Returns dict with keys: model, input_tokens, output_tokens, total_tokens,
    finish_reason, tool_call_count. Any field missing is None.
    """
    out = {"model": None, "input_tokens": None, "output_tokens": None,
           "total_tokens": None, "finish_reason": None, "tool_call_count": None}
    try:
        rj = json.loads(req_raw) if req_raw else {}
    except Exception:
        rj = {}
    try:
        sj = json.loads(resp_raw) if resp_raw else {}
    except Exception:
        sj = {}
    # Model: response.model > request.model
    out["model"] = sj.get("model") or rj.get("model")
    # Usage: support both OpenAI shape and Anthropic-passthrough shape
    usage = sj.get("usage") or {}
    if isinstance(usage, dict):
        out["input_tokens"]  = usage.get("prompt_tokens") or usage.get("input_tokens")
        out["output_tokens"] = usage.get("completion_tokens") or usage.get("output_tokens")
        out["total_tokens"]  = usage.get("total_tokens")
        if out["total_tokens"] is None and (out["input_tokens"] or out["output_tokens"]):
            out["total_tokens"] = (out["input_tokens"] or 0) + (out["output_tokens"] or 0)
    # Finish reason + tool call count from first choice
    choices = sj.get("choices")
    if isinstance(choices, list) and choices:
        c0 = choices[0] or {}
        out["finish_reason"] = c0.get("finish_reason")
        msg = c0.get("message") or {}
        tcs = msg.get("tool_calls") or []
        if isinstance(tcs, list):
            out["tool_call_count"] = len(tcs)
    return out

if gw_log_rows:
    values = []
    for r in gw_log_rows:
        rid = r.request_id
        src = r.source_table
        if not rid or not src:
            continue
        parsed = _parse_gateway_payload(r.request_payload or "", r.response_payload or "")
        values.append((
            rid, src,
            r.client_request_id,
            r.request_time,
            int(r.status_code) if r.status_code is not None else None,
            int(r.execution_ms) if r.execution_ms is not None else None,
            r.request_payload,
            r.response_payload,
            int(r.request_size_bytes) if r.request_size_bytes is not None else None,
            int(r.response_size_bytes) if r.response_size_bytes is not None else None,
            r.served_entity_id,
            r.requester,
            parsed["model"],
            parsed["input_tokens"],
            parsed["output_tokens"],
            parsed["total_tokens"],
            parsed["finish_reason"],
            parsed["tool_call_count"],
            now,
        ))
    if values:
        with obs_conn.cursor() as cur:
            execute_values(
                cur,
                """INSERT INTO gateway_inference_logs
                   (request_id, source_table, client_request_id, request_time,
                    status_code, execution_ms, request_payload, response_payload,
                    request_size_bytes, response_size_bytes, served_entity_id,
                    requester, model, input_tokens, output_tokens, total_tokens,
                    finish_reason, tool_call_count, last_synced)
                   VALUES %s
                   ON CONFLICT (source_table, request_id) DO UPDATE SET
                       client_request_id = EXCLUDED.client_request_id,
                       request_time = EXCLUDED.request_time,
                       status_code = EXCLUDED.status_code,
                       execution_ms = EXCLUDED.execution_ms,
                       request_payload = EXCLUDED.request_payload,
                       response_payload = EXCLUDED.response_payload,
                       request_size_bytes = EXCLUDED.request_size_bytes,
                       response_size_bytes = EXCLUDED.response_size_bytes,
                       served_entity_id = EXCLUDED.served_entity_id,
                       requester = EXCLUDED.requester,
                       model = EXCLUDED.model,
                       input_tokens = EXCLUDED.input_tokens,
                       output_tokens = EXCLUDED.output_tokens,
                       total_tokens = EXCLUDED.total_tokens,
                       finish_reason = EXCLUDED.finish_reason,
                       tool_call_count = EXCLUDED.tool_call_count,
                       last_synced = EXCLUDED.last_synced""",
                values, page_size=200,
            )
            obs_conn.commit()
            gw_log_count = len(values)
print(f"✅ Upserted {gw_log_count} gateway inference-log rows (from Delta, payload-enriched)")

# COMMAND ----------

obs_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 3: Sync Knowledge Bases (Delta → Lakebase)

# COMMAND ----------

VS_EP_TABLE = f"{CATALOG}.{SCHEMA}.vector_search_endpoints"
VS_IDX_TABLE = f"{CATALOG}.{SCHEMA}.vector_search_indexes"
LB_INST_TABLE = f"{CATALOG}.{SCHEMA}.lakebase_instances"

kb_conn = get_lakebase_connection()
vs_ep_count = 0
vs_idx_count = 0
lb_inst_count = 0

# Ensure tables
with kb_conn.cursor() as cur:
    for ddl in [
        """CREATE TABLE IF NOT EXISTS vector_search_endpoints (
            endpoint_name TEXT PRIMARY KEY, endpoint_id TEXT, status TEXT,
            endpoint_type TEXT, num_indexes INT DEFAULT 0, creator TEXT,
            workspace_id TEXT, created_at TIMESTAMP,
            last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW())""",
        """CREATE TABLE IF NOT EXISTS vector_search_indexes (
            index_name TEXT NOT NULL, endpoint_name TEXT NOT NULL,
            index_type TEXT, primary_key TEXT, creator TEXT, workspace_id TEXT,
            detailed_state TEXT, indexed_row_count INT DEFAULT 0,
            ready BOOLEAN DEFAULT FALSE, status_message TEXT,
            source_table TEXT, embedding_model TEXT, pipeline_type TEXT,
            last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (endpoint_name, index_name))""",
        """CREATE TABLE IF NOT EXISTS vector_search_health_history (
            endpoint_name TEXT NOT NULL, status TEXT NOT NULL,
            num_indexes INT DEFAULT 0,
            recorded_at TIMESTAMP WITH TIME ZONE DEFAULT NOW())""",
        """CREATE TABLE IF NOT EXISTS lakebase_instances (
            instance_name TEXT PRIMARY KEY, instance_id TEXT, state TEXT,
            capacity TEXT, pg_version TEXT, read_write_dns TEXT, read_only_dns TEXT,
            creator TEXT, created_at TEXT,
            last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW())""",
    ]:
        try:
            cur.execute(ddl)
        except Exception as e:
            print(f"  DDL warning: {e}")
    kb_conn.commit()
print("✅ Knowledge bases tables ensured")

# COMMAND ----------

# Sync VS endpoints
print(f"▸ Syncing Vector Search endpoints from Delta: {VS_EP_TABLE} ...")
try:
    with kb_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE vector_search_endpoints")
        cur.execute("TRUNCATE TABLE vector_search_indexes")
        # Don't truncate health_history — it's append-only for uptime tracking
        cur.execute("TRUNCATE TABLE lakebase_instances")
        kb_conn.commit()
    ep_df = spark.read.table(VS_EP_TABLE)
    ep_rows = ep_df.collect()
    if ep_rows:
        with kb_conn.cursor() as cur:
            for r in ep_rows:
                cur.execute(
                    """INSERT INTO vector_search_endpoints
                       (endpoint_name, endpoint_id, status, endpoint_type, num_indexes, creator, created_at, last_synced)
                       VALUES (%s, %s, %s, %s, %s, %s, NULL, NOW())
                       ON CONFLICT (endpoint_name) DO UPDATE SET
                           status = EXCLUDED.status, endpoint_type = EXCLUDED.endpoint_type,
                           num_indexes = EXCLUDED.num_indexes, last_synced = NOW()""",
                    (r.endpoint_name, r.endpoint_id, r.status, r.endpoint_type,
                     r.num_indexes, r.creator))
                # Record health history
                cur.execute(
                    "INSERT INTO vector_search_health_history (endpoint_name, status, num_indexes) VALUES (%s, %s, %s)",
                    (r.endpoint_name, r.status, r.num_indexes))
                vs_ep_count += 1
            kb_conn.commit()
    print(f"  ✅ {vs_ep_count} endpoints synced")
except Exception as exc:
    print(f"  ⚠️  VS endpoints sync failed: {exc}")

# COMMAND ----------

# Sync VS indexes
print(f"▸ Syncing Vector Search indexes from Delta: {VS_IDX_TABLE} ...")
try:
    idx_df = spark.read.table(VS_IDX_TABLE)
    idx_rows = idx_df.collect()
    if idx_rows:
        with kb_conn.cursor() as cur:
            for r in idx_rows:
                cur.execute(
                    """INSERT INTO vector_search_indexes
                       (index_name, endpoint_name, index_type, primary_key, creator,
                        detailed_state, indexed_row_count, ready, status_message,
                        source_table, embedding_model, pipeline_type, last_synced)
                       VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW())
                       ON CONFLICT (endpoint_name, index_name) DO UPDATE SET
                           index_type = EXCLUDED.index_type, detailed_state = EXCLUDED.detailed_state,
                           indexed_row_count = EXCLUDED.indexed_row_count, ready = EXCLUDED.ready,
                           status_message = EXCLUDED.status_message, source_table = EXCLUDED.source_table,
                           embedding_model = EXCLUDED.embedding_model, pipeline_type = EXCLUDED.pipeline_type,
                           last_synced = NOW()""",
                    (r.index_name, r.endpoint_name, r.index_type, r.primary_key,
                     r.creator, r.detailed_state, r.indexed_row_count, r.ready,
                     r.status_message, r.source_table, r.embedding_model, r.pipeline_type))
                vs_idx_count += 1
            kb_conn.commit()
    print(f"  ✅ {vs_idx_count} indexes synced")
except Exception as exc:
    print(f"  ⚠️  VS indexes sync failed: {exc}")

# COMMAND ----------

# Sync Lakebase instances
print(f"▸ Syncing Lakebase instances from Delta: {LB_INST_TABLE} ...")
try:
    lb_df = spark.read.table(LB_INST_TABLE)
    lb_rows = lb_df.collect()
    if lb_rows:
        with kb_conn.cursor() as cur:
            for r in lb_rows:
                cur.execute(
                    """INSERT INTO lakebase_instances
                       (instance_name, instance_id, state, capacity, pg_version,
                        read_write_dns, read_only_dns, creator, created_at, last_synced)
                       VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, NOW())
                       ON CONFLICT (instance_name) DO UPDATE SET
                           state = EXCLUDED.state, capacity = EXCLUDED.capacity,
                           pg_version = EXCLUDED.pg_version, last_synced = NOW()""",
                    (r.instance_name, r.instance_id, r.state, r.capacity,
                     r.pg_version, r.read_write_dns, r.read_only_dns,
                     r.creator, r.created_at))
                lb_inst_count += 1
            kb_conn.commit()
    print(f"  ✅ {lb_inst_count} instances synced")
except Exception as exc:
    print(f"  ⚠️  Lakebase instances sync failed: {exc}")

# COMMAND ----------

kb_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 4: Sync KB Billing (Delta → Lakebase)

# COMMAND ----------

KB_BILLING_DELTA = f"{CATALOG}.{SCHEMA}.kb_billing_daily"

kb_billing_conn = get_lakebase_connection()
vs_billing_count = 0
lb_billing_count = 0

# Ensure billing cache table
with kb_billing_conn.cursor() as cur:
    for ddl in [
        """CREATE TABLE IF NOT EXISTS kb_billing_daily (
            usage_date          DATE NOT NULL,
            product             TEXT NOT NULL,
            workspace_id        TEXT NOT NULL,
            endpoint_name       TEXT DEFAULT '',
            workload_type       TEXT DEFAULT 'other',
            total_dbus          NUMERIC(18,4) DEFAULT 0,
            total_cost_usd      NUMERIC(18,4) DEFAULT 0,
            last_synced         TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, product, workspace_id, endpoint_name, workload_type)
        )""",
        "CREATE INDEX IF NOT EXISTS idx_kbd_product ON kb_billing_daily (product)",
        "CREATE INDEX IF NOT EXISTS idx_kbd_ws ON kb_billing_daily (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_kbd_date ON kb_billing_daily (usage_date DESC)",
    ]:
        try:
            cur.execute(ddl)
        except Exception as e:
            print(f"  DDL warning: {e}")
    kb_billing_conn.commit()

# Sync from Delta (populated by 07_discover_kb_billing task)
print(f"▸ Syncing KB billing from Delta: {KB_BILLING_DELTA} ...")
try:
    with kb_billing_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE kb_billing_daily")
        kb_billing_conn.commit()
    kb_df = spark.read.table(KB_BILLING_DELTA)
    kb_rows = kb_df.collect()
    if kb_rows:
        values = [(r.usage_date, r.product, r.workspace_id,
                   r.endpoint_name or "", r.workload_type or "other",
                   float(r.total_dbus or 0), float(r.total_cost_usd or 0))
                  for r in kb_rows]
        vs_billing_count = sum(1 for r in kb_rows if r.product == "VECTOR_SEARCH")
        lb_billing_count = len(kb_rows) - vs_billing_count
        with kb_billing_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO kb_billing_daily
                   (usage_date, product, workspace_id, endpoint_name, workload_type, total_dbus, total_cost_usd, last_synced)
                   VALUES %s
                   ON CONFLICT (usage_date, product, workspace_id, endpoint_name, workload_type) DO UPDATE SET
                       total_dbus = EXCLUDED.total_dbus, total_cost_usd = EXCLUDED.total_cost_usd,
                       last_synced = NOW()""",
                [(v[0], v[1], v[2], v[3], v[4], v[5], v[6], now) for v in values],
                page_size=500)
            kb_billing_conn.commit()
    print(f"  ✅ Synced {vs_billing_count} VS + {lb_billing_count} Lakebase billing rows")
except Exception as exc:
    print(f"  ⚠️  KB billing sync failed: {exc}")

kb_billing_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 5: Sync User Analytics (Delta → Lakebase)

# COMMAND ----------

UA_DAILY_TABLE = f"{CATALOG}.{SCHEMA}.user_analytics_daily"
UA_HEATMAP_TABLE = f"{CATALOG}.{SCHEMA}.user_analytics_heatmap"

ua_conn = get_lakebase_connection()
ua_daily_count = 0
ua_heatmap_count = 0

# Ensure cache tables
with ua_conn.cursor() as cur:
    for ddl in [
        """CREATE TABLE IF NOT EXISTS user_analytics_daily (
            usage_date      DATE NOT NULL,
            requester       TEXT NOT NULL,
            endpoint_name   TEXT NOT NULL,
            request_count   BIGINT DEFAULT 0,
            total_tokens    BIGINT DEFAULT 0,
            last_synced     TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, requester, endpoint_name)
        )""",
        "CREATE INDEX IF NOT EXISTS idx_uad_date ON user_analytics_daily (usage_date DESC)",
        "CREATE INDEX IF NOT EXISTS idx_uad_user ON user_analytics_daily (requester)",
        """CREATE TABLE IF NOT EXISTS user_analytics_heatmap (
            dow             INT NOT NULL,
            hour            INT NOT NULL,
            request_count   BIGINT DEFAULT 0,
            period_days     INT DEFAULT 30,
            last_synced     TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (dow, hour, period_days)
        )""",
    ]:
        try:
            cur.execute(ddl)
        except Exception as e:
            print(f"  DDL warning: {e}")
    ua_conn.commit()

# COMMAND ----------

# Sync daily activity from Delta
print(f"▸ Syncing user analytics daily from Delta: {UA_DAILY_TABLE} ...")
try:
    with ua_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE user_analytics_daily")
        cur.execute("TRUNCATE TABLE user_analytics_heatmap")
        ua_conn.commit()
    ua_df = spark.read.table(UA_DAILY_TABLE)
    ua_rows = ua_df.collect()
    if ua_rows:
        values = [(r.usage_date, r.requester, r.endpoint_name, r.request_count, r.total_tokens, now)
                  for r in ua_rows]
        ua_daily_count = len(values)
        with ua_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO user_analytics_daily
                   (usage_date, requester, endpoint_name, request_count, total_tokens, last_synced)
                   VALUES %s""",
                values, page_size=500)
            ua_conn.commit()
    print(f"  ✅ {ua_daily_count} daily rows synced")
except Exception as exc:
    print(f"  ⚠️  UA daily sync failed: {exc}")

# COMMAND ----------

# Sync heatmap from Delta
print(f"▸ Syncing heatmap from Delta: {UA_HEATMAP_TABLE} ...")
try:
    hm_df = spark.read.table(UA_HEATMAP_TABLE)
    hm_rows = hm_df.collect()
    if hm_rows:
        values = [(r.dow, r.hour, r.request_count, r.period_days, now) for r in hm_rows]
        ua_heatmap_count = len(values)
        with ua_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO user_analytics_heatmap (dow, hour, request_count, period_days, last_synced)
                   VALUES %s""",
                values, page_size=500)
            ua_conn.commit()
    print(f"  ✅ {ua_heatmap_count} heatmap rows synced")
except Exception as exc:
    print(f"  ⚠️  UA heatmap sync failed: {exc}")

# COMMAND ----------

ua_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 6: Sync Gateway Usage (Delta → Lakebase)

# COMMAND ----------

GW_DAILY_TABLE = f"{CATALOG}.{SCHEMA}.gateway_usage_daily"
GW_HOURLY_TABLE = f"{CATALOG}.{SCHEMA}.gateway_usage_hourly"

gw_conn = get_lakebase_connection()
gw_daily_count = 0
gw_hourly_count = 0

with gw_conn.cursor() as cur:
    for ddl in [
        """CREATE TABLE IF NOT EXISTS gateway_usage_daily (
            usage_date      DATE NOT NULL,
            endpoint_name   TEXT NOT NULL,
            requester       TEXT DEFAULT '',
            request_count   BIGINT DEFAULT 0,
            input_tokens    BIGINT DEFAULT 0,
            output_tokens   BIGINT DEFAULT 0,
            error_count     BIGINT DEFAULT 0,
            rate_limited_count BIGINT DEFAULT 0,
            last_synced     TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, endpoint_name, requester)
        )""",
        "CREATE INDEX IF NOT EXISTS idx_gud_date ON gateway_usage_daily (usage_date DESC)",
        "CREATE INDEX IF NOT EXISTS idx_gud_ep ON gateway_usage_daily (endpoint_name)",
        # ADD COLUMN for tables that pre-existed without rate_limited_count
        "ALTER TABLE gateway_usage_daily  ADD COLUMN IF NOT EXISTS rate_limited_count BIGINT DEFAULT 0",
        "ALTER TABLE gateway_usage_hourly ADD COLUMN IF NOT EXISTS rate_limited_count BIGINT DEFAULT 0",
        """CREATE TABLE IF NOT EXISTS gateway_usage_hourly (
            hour            TEXT NOT NULL,
            endpoint_name   TEXT DEFAULT '',
            request_count   BIGINT DEFAULT 0,
            input_tokens    BIGINT DEFAULT 0,
            output_tokens   BIGINT DEFAULT 0,
            error_count     BIGINT DEFAULT 0,
            rate_limited_count BIGINT DEFAULT 0,
            last_synced     TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (hour, endpoint_name)
        )""",
    ]:
        try:
            cur.execute(ddl)
        except Exception as e:
            print(f"  DDL warning: {e}")
    gw_conn.commit()

# COMMAND ----------

# Sync daily
print(f"▸ Syncing gateway daily usage from Delta: {GW_DAILY_TABLE} ...")
try:
    with gw_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE gateway_usage_daily")
        cur.execute("TRUNCATE TABLE gateway_usage_hourly")
        gw_conn.commit()
    gw_df = spark.read.table(GW_DAILY_TABLE)
    gw_rows = gw_df.collect()
    if gw_rows:
        values = [(r.usage_date, r.endpoint_name, r.requester or "",
                   r.request_count, r.input_tokens, r.output_tokens, r.error_count,
                   getattr(r, "rate_limited_count", 0) or 0, now)
                  for r in gw_rows]
        gw_daily_count = len(values)
        with gw_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO gateway_usage_daily
                   (usage_date, endpoint_name, requester, request_count, input_tokens, output_tokens, error_count, rate_limited_count, last_synced)
                   VALUES %s""",
                values, page_size=500)
            gw_conn.commit()
    print(f"  ✅ {gw_daily_count} daily rows synced")
except Exception as exc:
    print(f"  ⚠️  Gateway daily sync failed: {exc}")

# COMMAND ----------

# Sync hourly
print(f"▸ Syncing gateway hourly usage from Delta: {GW_HOURLY_TABLE} ...")
try:
    gh_df = spark.read.table(GW_HOURLY_TABLE)
    gh_rows = gh_df.collect()
    if gh_rows:
        values = [(r.hour, r.endpoint_name or "", r.request_count,
                   r.input_tokens, r.output_tokens, r.error_count,
                   getattr(r, "rate_limited_count", 0) or 0, now)
                  for r in gh_rows]
        gw_hourly_count = len(values)
        with gw_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO gateway_usage_hourly
                   (hour, endpoint_name, request_count, input_tokens, output_tokens, error_count, rate_limited_count, last_synced)
                   VALUES %s""",
                values, page_size=500)
            gw_conn.commit()
    print(f"  ✅ {gw_hourly_count} hourly rows synced")
except Exception as exc:
    print(f"  ⚠️  Gateway hourly sync failed: {exc}")

# COMMAND ----------

gw_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Sync Unity AI Gateway (v2) usage summary (Delta → Lakebase)
# MAGIC Mirrors `uag_usage_summary` from `11_discover_ai_gateway_usage`.

# COMMAND ----------

UAG_TABLE = f"{CATALOG}.{SCHEMA}.uag_usage_summary"
uag_conn = get_lakebase_connection()
uag_count = 0

with uag_conn.cursor() as cur:
    cur.execute(
        """CREATE TABLE IF NOT EXISTS uag_usage_summary (
            endpoint_name         TEXT NOT NULL,
            request_count         BIGINT DEFAULT 0,
            input_tokens          BIGINT DEFAULT 0,
            output_tokens         BIGINT DEFAULT 0,
            cache_read_tokens     BIGINT DEFAULT 0,
            cache_creation_tokens BIGINT DEFAULT 0,
            p50_latency_ms        BIGINT DEFAULT 0,
            p95_latency_ms        BIGINT DEFAULT 0,
            p95_ttfb_ms           BIGINT DEFAULT 0,
            error_count           BIGINT DEFAULT 0,
            unique_users          BIGINT DEFAULT 0,
            max_event_time        TEXT,
            last_synced           TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (endpoint_name)
        )""")
    uag_conn.commit()

print(f"▸ Syncing {UAG_TABLE} → uag_usage_summary ...")
try:
    with uag_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE uag_usage_summary")
        uag_conn.commit()
    uag_rows = spark.read.table(UAG_TABLE).collect()
    if uag_rows:
        values = [(r.endpoint_name, int(r.request_count or 0), int(r.input_tokens or 0),
                   int(r.output_tokens or 0), int(r.cache_read_tokens or 0), int(r.cache_creation_tokens or 0),
                   int(r.p50_latency_ms or 0), int(r.p95_latency_ms or 0), int(r.p95_ttfb_ms or 0),
                   int(r.error_count or 0), int(r.unique_users or 0), r.max_event_time, now)
                  for r in uag_rows]
        uag_count = len(values)
        with uag_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO uag_usage_summary
                   (endpoint_name, request_count, input_tokens, output_tokens, cache_read_tokens,
                    cache_creation_tokens, p50_latency_ms, p95_latency_ms, p95_ttfb_ms,
                    error_count, unique_users, max_event_time, last_synced)
                   VALUES %s
                   ON CONFLICT (endpoint_name) DO UPDATE SET
                       request_count = EXCLUDED.request_count,
                       input_tokens = EXCLUDED.input_tokens,
                       output_tokens = EXCLUDED.output_tokens,
                       cache_read_tokens = EXCLUDED.cache_read_tokens,
                       cache_creation_tokens = EXCLUDED.cache_creation_tokens,
                       p50_latency_ms = EXCLUDED.p50_latency_ms,
                       p95_latency_ms = EXCLUDED.p95_latency_ms,
                       p95_ttfb_ms = EXCLUDED.p95_ttfb_ms,
                       error_count = EXCLUDED.error_count,
                       unique_users = EXCLUDED.unique_users,
                       max_event_time = EXCLUDED.max_event_time,
                       last_synced = NOW()""",
                values, page_size=500)
            uag_conn.commit()
    print(f"  ✅ {uag_count} UAG usage rows synced")
except Exception as exc:
    print(f"  ⚠️  uag_usage_summary sync failed: {exc}")

# Sync uag_usage_breakdown (agent-vs-human / by model / by api_type)
UAG_BREAKDOWN_TABLE = f"{CATALOG}.{SCHEMA}.uag_usage_breakdown"
uag_bd_count = 0
with uag_conn.cursor() as cur:
    cur.execute(
        """CREATE TABLE IF NOT EXISTS uag_usage_breakdown (
            dimension     TEXT NOT NULL,
            key           TEXT NOT NULL,
            request_count BIGINT DEFAULT 0,
            input_tokens  BIGINT DEFAULT 0,
            output_tokens BIGINT DEFAULT 0,
            cached_tokens BIGINT DEFAULT 0,
            last_synced   TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (dimension, key)
        )""")
    uag_conn.commit()
print(f"▸ Syncing {UAG_BREAKDOWN_TABLE} → uag_usage_breakdown ...")
try:
    with uag_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE uag_usage_breakdown")
        uag_conn.commit()
    bd_rows = spark.read.table(UAG_BREAKDOWN_TABLE).collect()
    if bd_rows:
        values = [(r.dimension, r.key, int(r.request_count or 0), int(r.input_tokens or 0),
                   int(r.output_tokens or 0), int(r.cached_tokens or 0), now)
                  for r in bd_rows]
        uag_bd_count = len(values)
        with uag_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO uag_usage_breakdown
                   (dimension, key, request_count, input_tokens, output_tokens, cached_tokens, last_synced)
                   VALUES %s
                   ON CONFLICT (dimension, key) DO UPDATE SET
                       request_count = EXCLUDED.request_count,
                       input_tokens = EXCLUDED.input_tokens,
                       output_tokens = EXCLUDED.output_tokens,
                       cached_tokens = EXCLUDED.cached_tokens,
                       last_synced = NOW()""",
                values, page_size=500)
            uag_conn.commit()
    print(f"  ✅ {uag_bd_count} UAG breakdown rows synced")
except Exception as exc:
    print(f"  ⚠️  uag_usage_breakdown sync failed: {exc}")

uag_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Phase 7: Sync Billing (Delta → Lakebase)
# MAGIC
# MAGIC Mirrors the four billing tables populated by `09_discover_billing`:
# MAGIC - `billing_serving_daily`
# MAGIC - `billing_token_daily`
# MAGIC - `billing_product_daily`
# MAGIC - `billing_user_endpoint_daily`
# MAGIC
# MAGIC Also stamps `billing_cache_meta` so the app's `get_cache_status()` surfaces fresh `last_refreshed` timestamps.

# COMMAND ----------

BSD_TABLE  = f"{CATALOG}.{SCHEMA}.billing_serving_daily"
BTD_TABLE  = f"{CATALOG}.{SCHEMA}.billing_token_daily"
BPD_TABLE  = f"{CATALOG}.{SCHEMA}.billing_product_daily"
BUED_TABLE = f"{CATALOG}.{SCHEMA}.billing_user_endpoint_daily"
BUCD_TABLE = f"{CATALOG}.{SCHEMA}.billing_user_cost_daily"

billing_conn = get_lakebase_connection()
bsd_count = 0
btd_count = 0
bpd_count = 0
bued_count = 0
bucd_count = 0

# Ensure billing tables (idempotent — also created by app's ensure_billing_tables on startup)
with billing_conn.cursor() as cur:
    for ddl in [
        """CREATE TABLE IF NOT EXISTS billing_serving_daily (
            usage_date    DATE          NOT NULL,
            workspace_id  TEXT          NOT NULL,
            endpoint_name TEXT          NOT NULL,
            sku_name      TEXT          NOT NULL DEFAULT '',
            total_dbus    NUMERIC(18,4) NOT NULL DEFAULT 0,
            total_cost_usd NUMERIC(18,4) NOT NULL DEFAULT 0,
            last_synced   TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, workspace_id, endpoint_name, sku_name)
        )""",
        """CREATE TABLE IF NOT EXISTS billing_token_daily (
            usage_date       DATE    NOT NULL,
            workspace_id     TEXT    NOT NULL,
            endpoint_name    TEXT    NOT NULL,
            request_count    BIGINT  NOT NULL DEFAULT 0,
            input_tokens     BIGINT  NOT NULL DEFAULT 0,
            output_tokens    BIGINT  NOT NULL DEFAULT 0,
            avg_input_tokens NUMERIC(12,2) NOT NULL DEFAULT 0,
            avg_output_tokens NUMERIC(12,2) NOT NULL DEFAULT 0,
            last_synced      TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, workspace_id, endpoint_name)
        )""",
        """CREATE TABLE IF NOT EXISTS billing_product_daily (
            usage_date              DATE          NOT NULL,
            workspace_id            TEXT          NOT NULL,
            billing_origin_product  TEXT          NOT NULL,
            total_dbus              NUMERIC(18,4) NOT NULL DEFAULT 0,
            total_cost_usd          NUMERIC(18,4) NOT NULL DEFAULT 0,
            last_synced             TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, workspace_id, billing_origin_product)
        )""",
        """CREATE TABLE IF NOT EXISTS billing_user_endpoint_daily (
            usage_date     DATE    NOT NULL,
            workspace_id   TEXT    NOT NULL,
            endpoint_name  TEXT    NOT NULL,
            user_identity  TEXT    NOT NULL DEFAULT '',
            request_count  BIGINT  NOT NULL DEFAULT 0,
            input_tokens   BIGINT  NOT NULL DEFAULT 0,
            output_tokens  BIGINT  NOT NULL DEFAULT 0,
            last_synced    TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, workspace_id, endpoint_name, user_identity)
        )""",
        """CREATE TABLE IF NOT EXISTS billing_user_cost_daily (
            usage_date     DATE          NOT NULL,
            workspace_id   TEXT          NOT NULL,
            endpoint_id    TEXT          NOT NULL DEFAULT '',
            endpoint_name  TEXT          NOT NULL DEFAULT '',
            run_by         TEXT          NOT NULL DEFAULT '',
            total_dbus     NUMERIC(18,4) NOT NULL DEFAULT 0,
            total_cost_usd NUMERIC(18,4) NOT NULL DEFAULT 0,
            last_synced    TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            PRIMARY KEY (usage_date, workspace_id, endpoint_id, run_by)
        )""",
        """CREATE TABLE IF NOT EXISTS billing_cache_meta (
            cache_key      TEXT PRIMARY KEY,
            last_refreshed TIMESTAMP WITH TIME ZONE,
            rows_loaded    INTEGER NOT NULL DEFAULT 0
        )""",
        "ALTER TABLE billing_cache_meta ADD COLUMN IF NOT EXISTS value_text TEXT",
        # ALTERs for tables that pre-existed without last_synced — CREATE TABLE
        # IF NOT EXISTS above is a no-op when the table already exists, so these
        # ADD COLUMN statements are needed to bring older schemas up to date.
        "ALTER TABLE billing_serving_daily        ADD COLUMN IF NOT EXISTS last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW()",
        "ALTER TABLE billing_token_daily          ADD COLUMN IF NOT EXISTS last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW()",
        "ALTER TABLE billing_product_daily        ADD COLUMN IF NOT EXISTS last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW()",
        "ALTER TABLE billing_user_endpoint_daily  ADD COLUMN IF NOT EXISTS last_synced TIMESTAMP WITH TIME ZONE DEFAULT NOW()",
        "CREATE INDEX IF NOT EXISTS idx_bsd_ws  ON billing_serving_daily  (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_btd_ws  ON billing_token_daily    (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_bpd_ws  ON billing_product_daily  (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_bued_ws ON billing_user_endpoint_daily (workspace_id)",
        "CREATE INDEX IF NOT EXISTS idx_bucd_ws ON billing_user_cost_daily (workspace_id)",
    ]:
        try:
            cur.execute(ddl)
        except Exception as e:
            print(f"  DDL warning: {e}")
    billing_conn.commit()


def _stamp_cache_meta(conn, cache_key: str, rows_loaded: int) -> None:
    """Update billing_cache_meta so the app's get_cache_status sees fresh timestamps."""
    try:
        with conn.cursor() as c:
            c.execute(
                """INSERT INTO billing_cache_meta (cache_key, last_refreshed, rows_loaded)
                   VALUES (%s, NOW(), %s)
                   ON CONFLICT (cache_key) DO UPDATE
                   SET last_refreshed = NOW(), rows_loaded = EXCLUDED.rows_loaded""",
                (cache_key, rows_loaded),
            )
            conn.commit()
    except Exception as exc:
        print(f"  ⚠️  cache_meta update failed for {cache_key}: {exc}")


# Sync billing_serving_daily
print(f"▸ Syncing {BSD_TABLE} → billing_serving_daily ...")
try:
    with billing_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE billing_serving_daily")
        billing_conn.commit()
    bsd_rows = spark.read.table(BSD_TABLE).collect()
    if bsd_rows:
        values = [(r.usage_date, r.workspace_id, r.endpoint_name, r.sku_name or "",
                   float(r.total_dbus or 0), float(r.total_cost_usd or 0))
                  for r in bsd_rows]
        bsd_count = len(values)
        with billing_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO billing_serving_daily
                   (usage_date, workspace_id, endpoint_name, sku_name, total_dbus, total_cost_usd, last_synced)
                   VALUES %s
                   ON CONFLICT (usage_date, workspace_id, endpoint_name, sku_name) DO UPDATE SET
                       total_dbus = EXCLUDED.total_dbus,
                       total_cost_usd = EXCLUDED.total_cost_usd,
                       last_synced = NOW()""",
                [(v[0], v[1], v[2], v[3], v[4], v[5], now) for v in values],
                page_size=500)
            billing_conn.commit()
    print(f"  ✅ {bsd_count} serving-cost rows synced")
    _stamp_cache_meta(billing_conn, "serving_daily", bsd_count)
except Exception as exc:
    print(f"  ⚠️  billing_serving_daily sync failed: {exc}")

# Sync billing_token_daily
print(f"▸ Syncing {BTD_TABLE} → billing_token_daily ...")
try:
    with billing_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE billing_token_daily")
        billing_conn.commit()
    btd_rows = spark.read.table(BTD_TABLE).collect()
    if btd_rows:
        values = [(r.usage_date, r.workspace_id, r.endpoint_name,
                   int(r.request_count or 0),
                   int(r.input_tokens or 0), int(r.output_tokens or 0),
                   float(r.avg_input_tokens or 0), float(r.avg_output_tokens or 0))
                  for r in btd_rows]
        btd_count = len(values)
        with billing_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO billing_token_daily
                   (usage_date, workspace_id, endpoint_name, request_count,
                    input_tokens, output_tokens, avg_input_tokens, avg_output_tokens, last_synced)
                   VALUES %s
                   ON CONFLICT (usage_date, workspace_id, endpoint_name) DO UPDATE SET
                       request_count = EXCLUDED.request_count,
                       input_tokens = EXCLUDED.input_tokens,
                       output_tokens = EXCLUDED.output_tokens,
                       avg_input_tokens = EXCLUDED.avg_input_tokens,
                       avg_output_tokens = EXCLUDED.avg_output_tokens,
                       last_synced = NOW()""",
                [(v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], now) for v in values],
                page_size=500)
            billing_conn.commit()
    print(f"  ✅ {btd_count} token-usage rows synced")
    _stamp_cache_meta(billing_conn, "token_daily", btd_count)
except Exception as exc:
    print(f"  ⚠️  billing_token_daily sync failed: {exc}")

# Sync billing_product_daily
print(f"▸ Syncing {BPD_TABLE} → billing_product_daily ...")
try:
    with billing_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE billing_product_daily")
        billing_conn.commit()
    bpd_rows = spark.read.table(BPD_TABLE).collect()
    if bpd_rows:
        values = [(r.usage_date, r.workspace_id, r.billing_origin_product,
                   float(r.total_dbus or 0), float(r.total_cost_usd or 0))
                  for r in bpd_rows]
        bpd_count = len(values)
        with billing_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO billing_product_daily
                   (usage_date, workspace_id, billing_origin_product, total_dbus, total_cost_usd, last_synced)
                   VALUES %s
                   ON CONFLICT (usage_date, workspace_id, billing_origin_product) DO UPDATE SET
                       total_dbus = EXCLUDED.total_dbus,
                       total_cost_usd = EXCLUDED.total_cost_usd,
                       last_synced = NOW()""",
                [(v[0], v[1], v[2], v[3], v[4], now) for v in values],
                page_size=500)
            billing_conn.commit()
    print(f"  ✅ {bpd_count} product-cost rows synced")
    _stamp_cache_meta(billing_conn, "product_daily", bpd_count)
except Exception as exc:
    print(f"  ⚠️  billing_product_daily sync failed: {exc}")

# Sync billing_user_endpoint_daily
print(f"▸ Syncing {BUED_TABLE} → billing_user_endpoint_daily ...")
try:
    with billing_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE billing_user_endpoint_daily")
        billing_conn.commit()
    bued_rows = spark.read.table(BUED_TABLE).collect()
    if bued_rows:
        values = [(r.usage_date, r.workspace_id, r.endpoint_name,
                   r.user_identity or "unknown",
                   int(r.request_count or 0),
                   int(r.input_tokens or 0), int(r.output_tokens or 0))
                  for r in bued_rows]
        bued_count = len(values)
        with billing_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO billing_user_endpoint_daily
                   (usage_date, workspace_id, endpoint_name, user_identity,
                    request_count, input_tokens, output_tokens, last_synced)
                   VALUES %s
                   ON CONFLICT (usage_date, workspace_id, endpoint_name, user_identity) DO UPDATE SET
                       request_count = EXCLUDED.request_count,
                       input_tokens = EXCLUDED.input_tokens,
                       output_tokens = EXCLUDED.output_tokens,
                       last_synced = NOW()""",
                [(v[0], v[1], v[2], v[3], v[4], v[5], v[6], now) for v in values],
                page_size=500)
            billing_conn.commit()
    print(f"  ✅ {bued_count} user-endpoint rows synced")
    _stamp_cache_meta(billing_conn, "user_endpoint_daily", bued_count)
except Exception as exc:
    print(f"  ⚠️  billing_user_endpoint_daily sync failed: {exc}")

# Sync billing_user_cost_daily (actual per-user $ — UAG v2 attribution)
print(f"▸ Syncing {BUCD_TABLE} → billing_user_cost_daily ...")
try:
    with billing_conn.cursor() as cur:
        cur.execute("TRUNCATE TABLE billing_user_cost_daily")
        billing_conn.commit()
    bucd_rows = spark.read.table(BUCD_TABLE).collect()
    if bucd_rows:
        values = [(r.usage_date, r.workspace_id, r.endpoint_id or "", r.endpoint_name or "",
                   r.run_by or "unknown",
                   float(r.total_dbus or 0), float(r.total_cost_usd or 0))
                  for r in bucd_rows]
        bucd_count = len(values)
        with billing_conn.cursor() as cur:
            execute_values(cur,
                """INSERT INTO billing_user_cost_daily
                   (usage_date, workspace_id, endpoint_id, endpoint_name, run_by,
                    total_dbus, total_cost_usd, last_synced)
                   VALUES %s
                   ON CONFLICT (usage_date, workspace_id, endpoint_id, run_by) DO UPDATE SET
                       endpoint_name = EXCLUDED.endpoint_name,
                       total_dbus = EXCLUDED.total_dbus,
                       total_cost_usd = EXCLUDED.total_cost_usd,
                       last_synced = NOW()""",
                [(v[0], v[1], v[2], v[3], v[4], v[5], v[6], now) for v in values],
                page_size=500)
            billing_conn.commit()
    print(f"  ✅ {bucd_count} user-cost rows synced")
    _stamp_cache_meta(billing_conn, "user_cost_daily", bucd_count)
except Exception as exc:
    print(f"  ⚠️  billing_user_cost_daily sync failed: {exc}")

billing_conn.close()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Final Summary

# COMMAND ----------

result = {
    "status": "success",
    **agent_result,
    "observability_experiments": exp_count,
    "observability_runs": run_count,
    "observability_traces": trace_count,
    "observability_trace_details": trace_detail_count,
    "traces_from_delta_default": delta_trace_count,
    "traces_from_delta_uc_otel": delta_uc_trace_count,
    "trace_details_from_delta_default": delta_detail_count,
    "trace_details_from_delta_uc_otel": delta_uc_detail_count,
    "gateway_inference_logs": gw_log_count,
    "gateway_inference_logs_from_delta": gw_log_delta_count,
    "vs_endpoints": vs_ep_count,
    "vs_indexes": vs_idx_count,
    "lakebase_instances": lb_inst_count,
    "kb_billing_vs": vs_billing_count,
    "kb_billing_lakebase": lb_billing_count,
    "ua_daily_rows": ua_daily_count,
    "ua_heatmap_rows": ua_heatmap_count,
    "gw_daily_rows": gw_daily_count,
    "gw_hourly_rows": gw_hourly_count,
    "uag_usage_rows": uag_count,
    "billing_serving_daily_rows": bsd_count,
    "billing_token_daily_rows": btd_count,
    "billing_product_daily_rows": bpd_count,
    "billing_user_endpoint_daily_rows": bued_count,
    "billing_user_cost_daily_rows": bucd_count,
    "synced_at": datetime.now(timezone.utc).isoformat(),
}
print(json.dumps(result, indent=2))
dbutils.notebook.exit(json.dumps(result))
