"""
Convexity SDK - Create a Demand-Forecasting Task

This sample demonstrates how to:
- Create a Task with the FORECASTING scenario (or update it if it already exists)
- Attach a single notebook code-cell that queries the project's DuckLake
  warehouse for monthly demand data and forecasts the next 3 months
- Create and start a Run, then poll until it completes
- Print the captured stdout output from the execution

The cell code receives ``CONVEXITY_API_KEY`` from the caller (injected as a
placeholder at build time) and reads ``CONVEXITY_BASE_URL`` and
``CONVEXITY_PROJECT_ID`` from environment variables that the backend injects
automatically when executing a run.

Prerequisites:
- A valid API key with WRITE or ADMIN permission
- The "Convexity SDK samples" org and "griffin it" project must exist,
  along with the "monthly_demand" dataset.
  Run the previous samples first:
    1. create_org_and_project.py
    2. create_databricks_connections.py
    3. create_datasets.py
"""

import os
import sys
import time

from convexity_api_client.models import CreateRunTaskPayload
from convexity_api_client.models.cell_type import CellType
from convexity_api_client.models.create_run_request import CreateRunRequest
from convexity_api_client.models.notebook_cell import NotebookCell
from convexity_api_client.models.notebook_document import NotebookDocument
from convexity_api_client.models.task_scenario import TaskScenario

from convexity_sdk import ConvexityClient

ORG_SLUG = "convexity-sdk-samples"
PROJECT_SLUG = "griffin-it"
POLL_INTERVAL = 3  # seconds between run 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 monthly_demand dataset exists ───────────────────────
ds_resp = client.datasets.list(project.id)
ds_by_name = {d.name: d for d in ds_resp.datasets}
if "monthly_demand" not in ds_by_name:
    print(
        'ERROR: Dataset "monthly_demand" not found.\nPlease run create_datasets.py first.',
        file=sys.stderr,
    )
    raise SystemExit(1)

monthly_demand_ds = ds_by_name["monthly_demand"]
print(f"  Found dataset: monthly_demand (id={monthly_demand_ds.id})")

# DuckLake tables are named ds_{dataset_id} with hyphens replaced by underscores
monthly_demand_table = f"ds_{monthly_demand_ds.id.replace('-', '_')}"

# ── Step 4: Create or update the forecasting task ───────────────────────
TASK_NAME = "Monthly Demand Forecast – Next 3 Months"

print("\nChecking for existing task …")
existing_tasks = client.tasks.list(project.id, scenario=TaskScenario.FORECASTING)
task = None
for t in existing_tasks.tasks:
    if t.name == TASK_NAME:
        task = t
        break

if task is not None:
    print(f"  Task already exists: {task.name} (id={task.id}) – will update")
else:
    print("Creating forecasting task …")
    task = client.tasks.create(
        name=TASK_NAME,
        project_id=project.id,
        scenario=TaskScenario.FORECASTING,
        description="Query monthly demand from the DuckLake warehouse and forecast the next 3 months using linear regression.",
    )
    print(f"  Created task: {task.name} (id={task.id})")

# ── Step 5: Build the notebook with a single code cell ──────────────────
#
# The cell code runs server-side. CONVEXITY_API_KEY is injected from the
# caller via placeholder substitution, while CONVEXITY_BASE_URL and
# CONVEXITY_PROJECT_ID come from environment variables that the backend
# injects automatically.
#
# The DuckLake table name is also baked in via placeholder.
#
CELL_CODE = """\
import os
import duckdb
import statistics

from convexity_sdk import ConvexityClient
# ── Read injected environment variables ────────────────────────────────
project_id = os.environ["CONVEXITY_PROJECT_ID"]
os.environ["CONVEXITY_API_KEY"] = "__CONVEXITY_API_KEY__"

# ── Fetch DuckLake connection info via the SDK ─────────────────────────
client = ConvexityClient()

dl_resp = client.datalake.get_info(project_id=project_id)
if dl_resp is None or not dl_resp.available or dl_resp.connection is None:
    raise RuntimeError("Datalake not available for this project.")

conn = dl_resp.connection
endpoint = (conn.s_3_endpoint or "").replace("https://", "").replace("http://", "")
ducklake_name = conn.ducklake_name or "datalake"

# ── Connect to the project DuckLake warehouse ──────────────────────────
db = duckdb.connect()
db.execute("INSTALL httpfs; LOAD httpfs;")
db.execute("INSTALL ducklake; LOAD ducklake;")

db.execute(f\"\"\"CREATE SECRET convexity_s3 (
    TYPE S3,
    KEY_ID '{conn.s_3_access_key_id}',
    SECRET '{conn.s_3_secret_access_key}',
    ENDPOINT '{endpoint}',
    USE_SSL true,
    URL_STYLE 'path',
    REGION '{conn.s_3_region or "auto"}'
);\"\"\")

db.execute(f"ATTACH 'ducklake:{conn.catalog_path}' AS {ducklake_name} (DATA_PATH '{conn.data_path}', READ_ONLY);")
db.execute(f"USE {ducklake_name};")

client.close()

# ── Query monthly demand, ordered by month ─────────────────────────────
MONTHLY_DEMAND_TABLE = "__MONTHLY_DEMAND_TABLE__"

rows = db.execute(f\"\"\"
    SELECT
        month,
        SUM(demand) AS total_demand
    FROM {MONTHLY_DEMAND_TABLE}
    GROUP BY month
    ORDER BY month ASC
\"\"\").fetchall()

if not rows:
    print("No monthly demand data found.")
else:
    months = [r[0] for r in rows]
    demands = [float(r[1]) for r in rows]

    print("=== Historical Monthly Demand ===")
    for m, d in zip(months, demands):
        print(f"  {str(m):>20}  |  {d:>12,.0f}")

    # ── Simple linear-regression forecast ──────────────────────────────
    n = len(demands)
    x = list(range(n))
    x_mean = statistics.mean(x)
    y_mean = statistics.mean(demands)

    numerator = sum((xi - x_mean) * (yi - y_mean) for xi, yi in zip(x, demands))
    denominator = sum((xi - x_mean) ** 2 for xi in x)
    slope = numerator / denominator if denominator else 0
    intercept = y_mean - slope * x_mean

    print("\\n=== 3-Month Demand Forecast (linear trend) ===")
    for i in range(1, 4):
        forecast_value = intercept + slope * (n - 1 + i)
        print(f"  Month +{i}  |  {forecast_value:>12,.0f}")

    print(f"\\nTrend slope: {slope:,.2f} units/month")

db.close()
"""

# Substitute placeholders
cell_code = CELL_CODE.replace("__MONTHLY_DEMAND_TABLE__", monthly_demand_table)
cell_code = cell_code.replace("__CONVEXITY_API_KEY__", os.environ["CONVEXITY_API_KEY"])

print("\nAttaching notebook cell to task …")
updated_task = client.tasks.update(
    task.id,
    notebook=NotebookDocument(
        cells=[
            NotebookCell(
                id="forecast-cell-1",
                type_=CellType.CODE,
                content=cell_code,
            ),
        ]
    ),
)

print("  Task updated with 1 code cell")

# ── Step 6: Create and start a run ──────────────────────────────────────
print("\nStarting run …")
run_resp = client.runs.create_and_start(
    body=CreateRunRequest(
        code=cell_code,
        task=CreateRunTaskPayload(
            id=updated_task.id,
            notebook=updated_task.notebook,
        ),
    )
)
if run_resp is None or not run_resp.success:
    msg = getattr(run_resp, "message", "unknown error") if run_resp else "no response"
    print(f"ERROR: Failed to create/start run – {msg}", file=sys.stderr)
    raise SystemExit(1)

run_id = run_resp.run_id
print(f"  Run created: {run_id}")

# ── Step 7: Poll until the run finishes ─────────────────────────────────
print("\nPolling for run completion …")
while True:
    run = client.runs.get_status(run_id)
    if run is None:
        print("ERROR: Could not fetch run status.", file=sys.stderr)
        raise SystemExit(1)

    status = run.status
    progress = getattr(run, "progress", None) or 0
    print(f"  status={status}  progress={progress:.0%}", end="\r")

    if status in ("completed", "failed", "cancelled"):
        print()  # newline after \r
        break

    time.sleep(POLL_INTERVAL)

# ── Step 8: Print the result ────────────────────────────────────────────
if status == "completed":
    print("\n── Run Output ──────────────────────────────────────────────")
    output = getattr(run, "output", None) or "(no output captured)"
    print(output)
    exec_ms = getattr(run, "execution_time_ms", None)
    if exec_ms is not None:
        print(f"\nExecution time: {exec_ms} ms")
elif status == "failed":
    error = getattr(run, "error_message", None) or "(no error message)"
    print(f"\nERROR: Run failed – {error}", file=sys.stderr)
    raise SystemExit(1)
else:
    print(f"\nRun ended with status: {status}")

# ── Cleanup ─────────────────────────────────────────────────────────────
client.close()
print("\nDone!")
