{#- Context variables: demo_title, catalog, warehouse_id, databricks_profile -#}
#!/usr/bin/env python3
"""{{ demo_title }} — Genie Space Setup

Creates a sample Genie Space with the gold_customer_summary and
silver_customers tables. Run locally (uses ~/.databrickscfg profile)
or from a Databricks notebook.

Usage:
    python scripts/setup-genie.py
    python scripts/setup-genie.py --profile MY_PROFILE
    python scripts/setup-genie.py --catalog my_catalog

This file is replaced by /plan + /build-next with real
industry-specific Genie configuration.
"""

from __future__ import annotations

import argparse
import os
import sys


def main() -> None:
    parser = argparse.ArgumentParser(description="Create a sample Genie Space")
    parser.add_argument(
        "--profile",
        default=os.environ.get(
            "DATABRICKS_CONFIG_PROFILE", "{{ databricks_profile or 'DEFAULT' }}"
        ),
        help="Databricks CLI profile (default: from .env or {{ databricks_profile or 'DEFAULT' }})",
    )
    parser.add_argument(
        "--catalog",
        default=os.environ.get("CATALOG", "{{ catalog }}"),
        help="Unity Catalog name (default: {{ catalog }})",
    )
    parser.add_argument(
        "--warehouse-id",
        default=os.environ.get("DATABRICKS_SQL_WAREHOUSE_ID", "{{ warehouse_id }}"),
        help="SQL Warehouse ID",
    )
    args = parser.parse_args()

    # Source .env if present
    env_file = os.path.join(os.path.dirname(__file__), "..", ".env")
    if os.path.exists(env_file):
        with open(env_file) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, _, value = line.partition("=")
                    os.environ.setdefault(key.strip(), value.strip())

    try:
        import json

        from databricks.sdk import WorkspaceClient
    except ImportError:
        print("ERROR: databricks-sdk not installed. Run: pip install databricks-sdk")
        sys.exit(1)

    w = WorkspaceClient(profile=args.profile)
    catalog = args.catalog
    warehouse_id = args.warehouse_id

    # The Genie API requires a serialized_space in protobuf-compatible JSON.
    # Version 2 format with data_sources.tables[].identifier for table refs.
    serialized_space = json.dumps(
        {
            "version": 2,
            "data_sources": {
                "tables": [
                    {"identifier": f"{catalog}.default.gold_customer_summary"},
                    {"identifier": f"{catalog}.default.silver_customers"},
                ]
            },
        }
    )

    try:
        space = w.genie.create_space(
            warehouse_id=warehouse_id,
            serialized_space=serialized_space,
            title="{{ demo_title }} — Sample Genie",
            description="Ask questions about the sample customer data",
        )
        host = (w.config.host or "").rstrip("/")
        print(f"Genie Space created: {space.space_id}")
        print(f"URL: {host}/genie/rooms/{space.space_id}")
    except Exception as exc:
        print(f"Warning: Genie Space creation failed: {exc}")
        print("You can create it manually from the Databricks UI")
        sys.exit(1)


if __name__ == "__main__":
    main()
