┌──────────────────────────────────────────────────────┐
│  1. WRITE BENCHMARKS                                  │
│     Generate domain-relevant benchmark questions      │
│     with expected SQL answers                         │
├──────────────────────────────────────────────────────┤
│  2. QUERY GENIE                                       │
│     Run benchmarks via Conversation API               │
│     (respect 12s rate limit between queries)          │
├──────────────────────────────────────────────────────┤
│  3. EVALUATE                                          │
│     Compare generated SQL to expected SQL             │
│     Measure accuracy + repeatability                  │
├──────────────────────────────────────────────────────┤
│  4. DIAGNOSE                                          │
│     Identify failure patterns:                        │
│     - Wrong asset selected (TVF vs MV vs Table)       │
│     - Wrong columns/aggregations                      │
│     - Missing routing context                         │
├──────────────────────────────────────────────────────┤
│  5. APPLY CONTROL LEVERS                              │
│     Fix issues using 6 levers (priority order)        │
│     Apply via BOTH API + repository (dual persist)    │
├──────────────────────────────────────────────────────┤
│  6. RE-TEST                                           │
│     Wait 30s → re-run failing questions               │
│     Measure improvement                               │
│     Loop back to step 4 if targets not met            │
├──────────────────────────────────────────────────────┤
│  7. DOCUMENT                                          │
│     Generate optimization report with metrics         │
└──────────────────────────────────────────────────────┘

benchmarks:
  - id: "cost_001"
    question: "What is total spend this month?"
    expected_sql: "SELECT MEASURE(total_cost) FROM ${catalog}.${schema}.mv_cost_analytics WHERE ..."
    expected_asset: "MV"
    category: "aggregation"

  - id: "cost_002"
    question: "Show top 10 costliest workspaces"
    expected_sql: "SELECT * FROM ${catalog}.${schema}.get_top_cost_contributors('7', 'workspace')"
    expected_asset: "TVF"
    category: "list"

import time
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

def run_genie_query(space_id: str, question: str, max_wait: int = 120) -> dict:
    """Execute a query against Genie and return SQL + status."""
    try:
        resp = w.genie.start_conversation(space_id=space_id, content=question)
        conversation_id = resp.conversation_id
        message_id = resp.message_id

        poll_interval = 3
        start = time.time()
        while time.time() - start < max_wait:
            time.sleep(poll_interval)
            msg = w.genie.get_message(
                space_id=space_id,
                conversation_id=conversation_id,
                message_id=message_id
            )
            status = str(msg.status) if hasattr(msg, 'status') else 'UNKNOWN'
            if any(s in status for s in ['COMPLETED', 'FAILED', 'CANCELLED']):
                break
            poll_interval = min(poll_interval + 1, 10)

        sql = None
        if hasattr(msg, 'attachments') and msg.attachments:
            for att in msg.attachments:
                if hasattr(att, 'query') and att.query:
                    sql = att.query.query if hasattr(att.query, 'query') else str(att.query)
        return {"status": status, "sql": sql}
    except Exception as e:
        return {"status": "ERROR", "sql": None, "error": str(e)}

time.sleep(12)  # MANDATORY between every Genie API call

def evaluate_accuracy(result: dict, expected: dict) -> dict:
    """Evaluate if Genie returned correct SQL."""
    generated_sql = (result.get("sql") or "").lower().strip()
    expected_sql = expected.get("expected_sql", "").lower().strip()

    # Check asset routing
    expected_asset = expected.get("expected_asset", "").upper()
    uses_mv = "mv_" in generated_sql or "measure(" in generated_sql
    uses_tvf = "get_" in generated_sql
    actual_asset = "MV" if uses_mv else ("TVF" if uses_tvf else "TABLE")

    return {
        "question": expected["question"],
        "sql_generated": result.get("status") == "COMPLETED" and bool(generated_sql),
        "correct_asset": actual_asset == expected_asset,
        "actual_asset": actual_asset,
        "expected_asset": expected_asset,
    }

import hashlib
from collections import Counter

def test_repeatability(space_id: str, question: str, iterations: int = 3) -> dict:
    """Test SQL consistency across multiple runs."""
    hashes = []
    for i in range(iterations):
        sql = run_genie_query(space_id, question).get("sql", "")
        sql_hash = hashlib.md5(sql.lower().encode()).hexdigest()[:8] if sql else "NONE"
        hashes.append(sql_hash)
        time.sleep(12)

    most_common_count = Counter(hashes).most_common(1)[0][1]
    return {
        "question": question,
        "repeatability_pct": (most_common_count / len(hashes)) * 100,
        "unique_variants": len(set(hashes)),
    }

# Re-test after optimization
time.sleep(30)  # Wait for propagation

for question in failing_questions:
    result = run_genie_query(space_id, question["question"])
    evaluation = evaluate_accuracy(result, question)
    print(f"  {'PASS' if evaluation['correct_asset'] else 'FAIL'}: {question['question']}")
    time.sleep(12)

# Run from Databricks notebook or local with SDK configured
python scripts/genie_optimizer.py --space-id <ID> --benchmarks golden-queries.yaml

python scripts/repeatability_tester.py --space-id <ID> --iterations 3

# WRONG: Unsorted → "data_sources.tables must be sorted"
payload = {"serialized_space": json.dumps(config)}

# CORRECT: Sort all arrays first
config = sort_genie_config(config)
payload = {"serialized_space": json.dumps(config)}

# WRONG: API-only update (lost on next deployment)
subprocess.run(["databricks", "api", "patch", ...])

# CORRECT: API + repository
subprocess.run(["databricks", "api", "patch", ...])         # Direct
with open("src/genie/config.json", "w") as f: json.dump(...)  # Repository

# WRONG: Rapid-fire → queries fail silently
for q in questions: run_genie_query(space_id, q)

# CORRECT: 12s between each query
for q in questions:
    run_genie_query(space_id, q)
    time.sleep(12)