"""
Convexity SDK - Create Datasets from Databricks Connections

This sample demonstrates how to:
- Verify that the organization, project, and connections exist
- Create SQL datasets linked to Databricks connections
- Trigger a refresh on each dataset and poll until completion

Prerequisites:
- A valid API key with WRITE or ADMIN permission
- The "Convexity SDK samples" org, "griffin it" project, and both
  Databricks connections ("canonical", "analytic") must already exist.
  Run `create_org_and_project.py` and `create_databricks_connections.py`
  first if they don't.
"""

import sys
import time

from convexity_api_client.models.dataset_config import DatasetConfig
from convexity_api_client.models.dataset_type import DatasetType
from convexity_api_client.models.refresh_dataset_payload import RefreshDatasetPayload
from convexity_api_client.models.refresh_mode import RefreshMode
from convexity_api_client.models.sql_dataset_config import SqlDatasetConfig
from convexity_api_client.models.worker_task_status import WorkerTaskStatus

from convexity_sdk import ConvexityClient

ORG_SLUG = "convexity-sdk-samples"
PROJECT_SLUG = "griffin-it"
POLL_INTERVAL = 5  # seconds between refresh status polls

# ── Initialise SDK client ───────────────────────────────────────────────
client = ConvexityClient()

# ── Step 1: Verify organization exists ──────────────────────────────────
org = client.organizations.get_by_slug(ORG_SLUG)
if org is None or not hasattr(org, "id"):
    print(
        f'ERROR: Organization "{ORG_SLUG}" not found.\nPlease run create_org_and_project.py first.',
        file=sys.stderr,
    )
    raise SystemExit(1)
print(f"Found organization: {org.name} (id={org.id})")

# ── Step 2: Verify project exists ──────────────────────────────────────
project = client.projects.get_by_slug(org.id, PROJECT_SLUG)
if project is None or not hasattr(project, "id"):
    print(
        f'ERROR: Project "{PROJECT_SLUG}" not found in org "{ORG_SLUG}".\nPlease run create_org_and_project.py first.',
        file=sys.stderr,
    )
    raise SystemExit(1)
print(f"Found project: {project.name} (id={project.id})")

# ── Step 3: Verify connections exist ────────────────────────────────────
conn_list = client.connections.list(project_id=project.id)
if conn_list is None or not hasattr(conn_list, "connections"):
    print(
        "ERROR: Could not list connections.\nPlease run create_databricks_connections.py first.",
        file=sys.stderr,
    )
    raise SystemExit(1)

conn_by_name = {c.name: c for c in conn_list.connections}

canonical_conn = conn_by_name.get("canonical")
analytic_conn = conn_by_name.get("analytic")

if canonical_conn is None or analytic_conn is None:
    missing = []
    if canonical_conn is None:
        missing.append("canonical")
    if analytic_conn is None:
        missing.append("analytic")
    print(
        f"ERROR: Connection(s) {', '.join(missing)} not found.\nPlease run create_databricks_connections.py first.",
        file=sys.stderr,
    )
    raise SystemExit(1)

print(f"  Found connection: canonical (id={canonical_conn.id})")
print(f"  Found connection: analytic  (id={analytic_conn.id})")

# ── Dataset definitions ─────────────────────────────────────────────────
DATASETS = [
    {
        "name": "family_monthly_forecast_12mo",
        "connection": canonical_conn,
        "sql": ("select forecast_units, CAST(month AS TIMESTAMP) AS month from analytics.family_monthly_forecast_12mo\n  where forecast_units != 0"),
    },
    {
        "name": "revenue_trend",
        "connection": analytic_conn,
        "sql": """\
WITH daily_summary AS (
  SELECT
    sh.txn_date as date,
    TRUNC(sh.txn_date, 'MONTH') as month,
    SUM(sh.unit_price * sh.txn_qty) as daily_revenue,
    SUM(sh.unit_cost * sh.txn_qty) as daily_cost,
    SUM((sh.unit_price - sh.unit_cost) * sh.txn_qty) as daily_profit,
    COUNT(DISTINCT sh.transaction_id) as daily_transaction_count,
    COUNT(DISTINCT sh.item_id) as daily_unique_items,
    SUM(sh.txn_qty) as daily_units_sold
  FROM canonical.stockhistory sh
  WHERE sh.txn_type = 'sales_invoice'
  GROUP BY sh.txn_date, TRUNC(sh.txn_date, 'MONTH')
),
monthly_summary AS (
  SELECT
    month,
    SUM(daily_revenue) as monthly_revenue,
    SUM(daily_cost) as monthly_cost,
    SUM(daily_profit) as monthly_profit,
    COUNT(*) as days_in_month,
    ROUND(SUM(daily_profit) / SUM(daily_revenue) * 100, 2) as monthly_margin_pct
  FROM daily_summary
  GROUP BY month
)
SELECT
  ds.date,
  ds.month,
  ds.daily_revenue,
  ds.daily_cost,
  ds.daily_profit,
  ROUND(ds.daily_profit / ds.daily_revenue * 100, 2) as daily_margin_pct,
  ds.daily_transaction_count,
  ds.daily_unique_items,
  ds.daily_units_sold,

  -- Monthly metrics
  ms.monthly_revenue,
  ms.monthly_cost,
  ms.monthly_profit,
  ms.monthly_margin_pct,
  ms.days_in_month,

  -- YTD cumulative
  SUM(ds.daily_revenue) OVER (ORDER BY ds.date) as ytd_revenue,
  SUM(ds.daily_cost) OVER (ORDER BY ds.date) as ytd_cost,
  SUM(ds.daily_profit) OVER (ORDER BY ds.date) as ytd_profit,

  -- Previous month comparison
  LAG(ms.monthly_revenue) OVER (ORDER BY ds.month) as prev_month_revenue,
  ROUND(
    (ms.monthly_revenue - LAG(ms.monthly_revenue) OVER (ORDER BY ds.month)) /
    LAG(ms.monthly_revenue) OVER (ORDER BY ds.month) * 100,
    2
  ) as mom_growth_pct

FROM daily_summary ds
LEFT JOIN monthly_summary ms ON ds.month = ms.month
ORDER BY ds.date ASC""",
    },
    {
        "name": "kpi",
        "connection": analytic_conn,
        "sql": """\
SELECT
  -- Revenue & Profit Metrics
  SUM(sh.unit_price * sh.txn_qty) as total_revenue,
  SUM(sh.unit_cost * sh.txn_qty) as total_cost,
  SUM((sh.unit_price - sh.unit_cost) * sh.txn_qty) as gross_profit,
  ROUND(
    SUM((sh.unit_price - sh.unit_cost) * sh.txn_qty) /
    SUM(sh.unit_price * sh.txn_qty) * 100,
    2
  ) as profit_margin_pct,
  COUNT(DISTINCT sh.transaction_id) as transaction_count,

  -- Inventory Metrics
  (SELECT COUNT(*) FROM canonical.items) as total_items,
  (SELECT SUM(CASE WHEN active = 'Y' THEN 1 ELSE 0 END) FROM canonical.items) as active_items,
  (SELECT SUM(CASE WHEN active = 'N' THEN 1 ELSE 0 END) FROM canonical.items) as inactive_items,
  (SELECT SUM(qty_on_hand * avg_cost) FROM canonical.items) as total_inventory_value,
  (SELECT SUM(qty_in_transit * avg_cost) FROM canonical.items) as in_transit_value,
  (SELECT ROUND(AVG(price), 2) FROM canonical.items) as avg_product_price

FROM canonical.stockhistory sh
WHERE sh.txn_type IN ('SALE', 'INVOICE')
  AND YEAR(sh.txn_date) = YEAR(CURRENT_DATE())""",
    },
    {
        "name": "item_category_mapping",
        "connection": canonical_conn,
        "sql": "Select * from analytics.item_category_mapping",
    },
    {
        "name": "monthly_demand",
        "connection": canonical_conn,
        "sql": """\
SELECT
  CAST(month AS TIMESTAMP) AS month,
  item_id,
  demand
FROM analytics.monthly_demand
ORDER by month""",
    },
]


# ── Helper: create dataset ──────────────────────────────────────────────
def create_dataset(name: str, connection, sql: str):
    """Create a SQL dataset with no row limit."""
    ds = client.datasets.create(
        project_id=project.id,
        name=name,
        type_=DatasetType.SQL,
        config=DatasetConfig(
            sql=SqlDatasetConfig(query=sql),
        ),
        connection_id=connection.id,
        refresh_mode=RefreshMode.MANUAL,
    )
    return ds


# ── Helper: refresh and wait ────────────────────────────────────────────
def refresh_and_wait(dataset):
    """Trigger an async refresh and poll until it completes or fails."""
    try:
        resp = client.datasets.refresh(dataset.id)
    except ValueError as exc:
        print(
            f"  ERROR: Failed to trigger refresh for '{dataset.name}' – {exc}",
            file=sys.stderr,
        )
        return

    print(f"  Refresh started (task_id={resp.id}) …", end="", flush=True)
    while True:
        time.sleep(POLL_INTERVAL)
        try:
            state = client.datasets.get_refresh_status(dataset.id, resp.id)
        except ValueError:
            print(" error polling status")
            return

        # print current status from payload
        payload = state.payload
        if not isinstance(payload, RefreshDatasetPayload):
            raise ValueError(f"Expected payload of type RefreshDatasetPayload, got {type(payload)}")
        print(f" {payload.progress_message}", end="", flush=True)
        if state.status == WorkerTaskStatus.COMPLETED:
            print(f" COMPLETED ({payload.row_count} rows)")
            return
        if state.status == WorkerTaskStatus.FAILED:
            error_msg = getattr(state, "error", "unknown error")
            print(f" FAILED — {error_msg}")
            return
        print(".", end="", flush=True)


# ── Step 4: Create all datasets ────────────────────────────────────────
print("\nCreating datasets …")
created_datasets = []
for defn in DATASETS:
    ds = create_dataset(defn["name"], defn["connection"], defn["sql"])
    print(f"  Created: {ds.name} (id={ds.id})")
    created_datasets.append(ds)

# ── Step 5: Refresh each dataset one by one ─────────────────────────────
print("\nRefreshing datasets …")
for ds in created_datasets:
    print(f"\n  [{ds.name}]")
    refresh_and_wait(ds)

print("\nDone!")
client.close()
