"""
Convexity SDK - Cancel a Pending Dataset Refresh Task

This sample demonstrates how to:
- Resolve a project by slug
- Start an async dataset refresh
- Cancel the pending refresh task
- Poll for the final task status

Prerequisites:
- A valid API key with WRITE or ADMIN permission
- The organization, project, and at least one dataset must already exist
"""

import os
import sys
import time

from convexity_api_client.models.refresh_dataset_payload import RefreshDatasetPayload
from convexity_api_client.models.worker_task_status import WorkerTaskStatus

from convexity_sdk import ConvexityClient

ORG_SLUG = os.getenv("CONVEXITY_ORG_SLUG", "convexity-sdk-samples")
PROJECT_SLUG = os.getenv("CONVEXITY_PROJECT_SLUG", "griffin-it")
DATASET_NAME = os.getenv("CONVEXITY_DATASET_NAME")
CANCEL_DELAY_SECONDS = float(os.getenv("CONVEXITY_CANCEL_DELAY", "1"))
POLL_INTERVAL = 2

client = ConvexityClient()
api = client._api_client

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.',
        file=sys.stderr,
    )
    raise SystemExit(1)

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}".',
        file=sys.stderr,
    )
    raise SystemExit(1)

print(f"Found project: {project.name} (id={project.id})")

list_resp = client.datasets.list(project.id)
if not list_resp.datasets:
    print("ERROR: No datasets found to refresh.", file=sys.stderr)
    raise SystemExit(1)

if DATASET_NAME:
    dataset = next((ds for ds in list_resp.datasets if ds.name == DATASET_NAME), None)
    if dataset is None:
        print(
            f'ERROR: Dataset "{DATASET_NAME}" not found in project.',
            file=sys.stderr,
        )
        raise SystemExit(1)
else:
    dataset = list_resp.datasets[0]

print(f"Using dataset: {dataset.name} (id={dataset.id})")

try:
    task = client.datasets.refresh(dataset.id)
except ValueError as exc:
    print(f"ERROR: Failed to trigger refresh: {exc}", file=sys.stderr)
    raise SystemExit(1)

print(f"Refresh started (task_id={task.id}). Waiting {CANCEL_DELAY_SECONDS}s before cancelling...")

time.sleep(CANCEL_DELAY_SECONDS)

try:
    cancelled = client.datasets.cancel_refresh(dataset.id, task.id)
except ValueError as exc:
    print(f"ERROR: Cancel request failed: {exc}", file=sys.stderr)
    raise SystemExit(1)

print(f"Cancel request sent. Status now: {cancelled.status}")

while True:
    time.sleep(POLL_INTERVAL)
    try:
        state = client.datasets.get_refresh_status(dataset.id, task.id)
    except ValueError as exc:
        print(f"ERROR: Failed to poll task status: {exc}", file=sys.stderr)
        raise SystemExit(1)

    payload = state.payload
    if isinstance(payload, RefreshDatasetPayload):
        message = payload.progress_message or ""
    else:
        message = ""

    print(f"Status: {state.status} {message}")

    if state.status in {WorkerTaskStatus.CANCELLED, WorkerTaskStatus.COMPLETED, WorkerTaskStatus.FAILED}:
        break

client.close()
