# ---
# jupyter:
#   jupytext:
#     formats: py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Training smoke test (`input-sample.md`)
#
# Run cells top to bottom in VS Code / Cursor (**Run Cell** on each `# %%` block).
#
# Steps mirror `input-sample.md`: optional DB delete + insert, `mc cp` to MinIO, Temporal `train_model`.
# Set `POSTGRES_*`, `TEMPORAL_*`, `TRAIN_TASK_QUEUE`, and configure the `mc` alias (default `suse`).

# %%
from __future__ import annotations

import os
import subprocess
import uuid
from datetime import datetime, timedelta
from pathlib import Path

import psycopg2
from dotenv import load_dotenv
from psycopg2.extras import Json
from temporalio import client
from dotenv import load_dotenv

load_dotenv()

# %%
# --- configuration (edit here or use `.env` at repo root) ---
PROJECT_ROOT = Path(__file__).resolve().parent.parent
load_dotenv(PROJECT_ROOT / '.env')

EXPERIMENT_RUN_ID = 1001
MINIO_MC_ALIAS = os.getenv('MINIO_MC_ALIAS', 'suse')
MINIO_BUCKET = os.getenv('MINIO_DEFAULT_BUCKET', 'model-training')
OBJECT_NAME = f'training-sample-dataset-{EXPERIMENT_RUN_ID}.csv'
LOCAL_CSV = PROJECT_ROOT / 'input_dataset.csv'

PG = {
    'host': os.getenv('POSTGRES_HOST'),
    'port': int(os.getenv('POSTGRES_PORT', '5432')),
    'user': os.getenv('POSTGRES_USER'),
    'password': os.getenv('POSTGRES_PASSWORD'),
    'dbname': os.getenv('POSTGRES_DBNAME'),
}

TEMPORAL_HOST = os.getenv('TEMPORAL_HOST')
TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE')
TRAIN_TASK_QUEUE = "train_model-single-queue"
TEMPORAL_TLS = os.getenv('TEMPORAL_USE_TLS', 'false').lower() in ('1', 'true', 'yes')


print(MINIO_MC_ALIAS, MINIO_BUCKET, OBJECT_NAME, LOCAL_CSV)
print(PG)
print(TEMPORAL_HOST, TEMPORAL_NAMESPACE, TRAIN_TASK_QUEUE, TEMPORAL_TLS)

# %%
# --- 1) database: delete previous row (same id), then insert `experiment_run` ---
# Primary key column is `id` (see `experiment_tracking` updates). `request_data` matches the SQL sample in `input-sample.md`.

now = datetime.utcnow()
request_data = {
    'experiment_run_id': EXPERIMENT_RUN_ID,
    'variable_columns': ['feature_a', 'feature_b'],
    'target_variable': 'target',
    'bucket_name': MINIO_BUCKET,
    'file_name': OBJECT_NAME,
    'line_separator': ',',
    'decimal_separator': '.',
    'train_size': 80,
    'shuffle': True,
    'random_state': 42,
    'model_name': 'test-runtime-linear-regression-model',
    'model_type': 'linear_regression',
    'model_id': EXPERIMENT_RUN_ID,
    'data_model_kwargs': {},
    'model_kwargs': {},
    'opt_params': {},
    'model_metadata': {'schemas': {'components': {'schemas': {}}}},
}

with psycopg2.connect(**PG) as conn:
    with conn.cursor() as cur:
        cur.execute('DELETE FROM public.experiment_run WHERE id = %s', (EXPERIMENT_RUN_ID,))
        cur.execute(
            """
            INSERT INTO public.experiment_run (
                id, experiment_name, run_name, username, status, error_message,
                created_at, updated_at, bucket_name, file_name, request_data, orchestrator_response_data
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            (
                EXPERIMENT_RUN_ID,
                'test-experiment-name',
                'test-run-name',
                'vitor.santos@aignosi.com.br',
                'ORCHESTRATOR_WAITING_PROC',
                None,
                now,
                now,
                MINIO_BUCKET,
                OBJECT_NAME,
                Json(request_data),
                None,
            ),
        )

# %%
# --- 2) MinIO: upload local CSV (requires `mc` CLI and alias configured) ---
subprocess.run(
    ['mc', 'cp', str(LOCAL_CSV), f'{MINIO_MC_ALIAS}/{MINIO_BUCKET}/{OBJECT_NAME}'],
    check=True,
)

# %%
# --- 3) Temporal: start `train_model` (flat payload; worker fills `model_metadata` in `load_model_metadata`) ---

if not TEMPORAL_HOST or not TEMPORAL_NAMESPACE or not TRAIN_TASK_QUEUE:
    raise RuntimeError('Set TEMPORAL_HOST, TEMPORAL_NAMESPACE, and TRAIN_TASK_QUEUE')
TH, TN, TQ = TEMPORAL_HOST, TEMPORAL_NAMESPACE, TRAIN_TASK_QUEUE

_workflow_input = {
    'experiment_run_id': EXPERIMENT_RUN_ID,
    'variable_columns': ['feature_a', 'feature_b'],
    'target_variable': 'target',
    'bucket_name': MINIO_BUCKET,
    'file_name': OBJECT_NAME,
    'line_separator': ',',
    'decimal_separator': '.',
    'train_size': 80,
    'shuffle': True,
    'random_state': 42,
    'model_name': 'test-runtime-linear-regression-model',
    'model_type': 'linear_regression',
    'model_id': EXPERIMENT_RUN_ID,
    'data_model_kwargs': {},
    'model_kwargs': {},
    'opt_params': {},
}

c = await client.Client.connect(
    target_host=TH,
    namespace=TN,
    tls=TEMPORAL_TLS,
)

# %%

wid = f'train-model-test-{uuid.uuid4()}'
result = await c.execute_workflow(  # type: ignore[call-overload]
    'train_model',
    _workflow_input,
    id=wid,
    task_queue=TQ,
    execution_timeout=timedelta(minutes=5),
    run_timeout=timedelta(minutes=5),
    task_timeout=timedelta(minutes=5),
)
print(wid)
print(result)

# %%