#!/usr/bin/env python3
"""Utility script to trigger the training workflow end-to-end for testing.

Steps performed:
1. Upload the CSV test dataset to MinIO using the configured `mc` alias.
2. Insert a new experiment_run record in Postgres and capture the generated ID.
3. Trigger the Temporal `train_model` workflow with the correct payload.

Prerequisites:
- `mc` CLI configured with alias defined in MINIO_ALIAS.
- PostgreSQL accessible with credentials in environment variables or defaults.
- Temporal server reachable without TLS on TEMPORAL_HOST / TEMPORAL_NAMESPACE.
- Python dependencies installed (see requirements.txt / requirements-dev.txt).
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import subprocess
import sys
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


# Carrega variáveis de ambiente do arquivo .env na raiz do projeto
PROJECT_ROOT = Path(__file__).resolve().parent.parent
ENV_PATH = PROJECT_ROOT / '.env'
if ENV_PATH.exists():
    load_dotenv(dotenv_path=ENV_PATH)


DEFAULT_CSV_PATH = Path('docs/test-model-data.csv')
TEST_SCENARIOS_DIR = PROJECT_ROOT / 'docs' / 'test-scenarios'
MINIO_ALIAS = 'suse'
MINIO_BUCKET = 'model-training'

POSTGRES_CONFIG = {
    'host': os.getenv('POSTGRES_HOST'),
    'port': os.getenv('POSTGRES_PORT'),
    '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 = os.getenv('TRAIN_TASK_QUEUE')
TEMPORAL_WORKFLOW = 'train_model'


def list_available_scenarios() -> list[str]:
    """List all available test scenario files."""
    if not TEST_SCENARIOS_DIR.exists():
        return []
    return sorted([f.stem for f in TEST_SCENARIOS_DIR.glob('*.json')])


def load_scenario(scenario_name: str) -> dict:
    """Load a test scenario from JSON file.

    Args:
        scenario_name: Name of the scenario (without .json extension)
                       or full path to a JSON file.

    Returns:
        Dictionary with scenario data.

    Raises:
        FileNotFoundError: If scenario file doesn't exist.
    """
    # Check if it's a full path
    scenario_path = Path(scenario_name)
    if scenario_path.suffix == '.json' and scenario_path.exists():
        with open(scenario_path) as f:
            return json.load(f)

    # Otherwise, look in the test-scenarios directory
    scenario_file = TEST_SCENARIOS_DIR / f'{scenario_name}.json'
    if not scenario_file.exists():
        available = list_available_scenarios()
        available_str = ', '.join(available) if available else 'none'
        raise FileNotFoundError(
            f"Scenario '{scenario_name}' not found at {scenario_file}.\n"
            f'Available scenarios: {available_str}'
        )

    with open(scenario_file) as f:
        return json.load(f)


def _ensure_source_file(path: Path) -> None:
    if not path.exists():
        raise FileNotFoundError(f'Test dataset not found at {path.resolve()}')


def upload_to_minio(source_path: Path) -> str:
    """Upload the CSV to MinIO using the mc CLI and return the object name."""
    _ensure_source_file(source_path)
    timestamp = datetime.utcnow().strftime('%Y%m%d-%H%M%S')
    object_name = f'test-model-data-{timestamp}.csv'
    target_uri = f'{MINIO_ALIAS}/{MINIO_BUCKET}/{object_name}'

    subprocess.run(  # noqa: S603
        ['mc', 'cp', str(source_path), target_uri],  # noqa: S607
        check=True,
    )
    return object_name


def insert_experiment_run(file_name: str, request_data: dict) -> int:
    """Insert experiment_run record and return the generated ID."""
    now = datetime.utcnow()
    payload = {
        **request_data,
        'fileName': file_name,
        'bucketName': MINIO_BUCKET,
    }

    insert_sql = """
        INSERT INTO experiment_run (
            experiment_name,
            username,
            status,
            created_at,
            updated_at,
            bucket_name,
            file_name,
            request_data
        )
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING id;
    """

    with psycopg2.connect(**POSTGRES_CONFIG) as conn:
        with conn.cursor() as cur:
            cur.execute(
                insert_sql,
                (
                    request_data['experimentName'],
                    request_data['username'],
                    'ORCHESTRATOR_WAITING_PROC',
                    now,
                    now,
                    MINIO_BUCKET,
                    file_name,
                    Json(payload),
                ),
            )
            experiment_run_id = cur.fetchone()[0]

    return experiment_run_id


def build_workflow_payload(
    experiment_run_id: int,
    file_name: str,
    request_data: dict,
) -> dict:
    """Convert camelCase request data to snake_case and enrich with runtime values."""
    return {
        'experiment_run_id': experiment_run_id,
        'experiment_name': request_data['experimentName'],
        'username': request_data['username'],
        'target_variable': request_data['targetVariable'],
        'variable_columns': request_data['variableColumns'],
        'lag_train': request_data['lagTrain'],
        'lag_val': request_data['lagVal'],
        'rem_static_win': request_data['remStaticWin'],
        'low_lim': request_data['lowLim'],
        'upp_lim': request_data['uppLim'],
        'window': request_data['window'],
        'use_scaler': request_data['useScaler'],
        'include_ar': request_data['includeAr'],
        'train_size': request_data['trainSize'],
        'shuffle': request_data['shuffle'],
        'bucket_name': MINIO_BUCKET,
        'file_name': file_name,
        'line_separator': request_data['lineSeparator'],
        'decimal_separator': request_data['decimalSeparator'],
        'removed_intervals': request_data['removedIntervals'],
        # New parameters
        'model_name': request_data.get('modelName', 'Linear Regression'),
        'degree': request_data.get('degree', 1),
        'interaction_only': request_data.get('interactionOnly', False),
        'nan_treatment': request_data.get('nanTreatment', 'drop'),
        'start_date': request_data.get('startDate'),
        'end_date': request_data.get('endDate'),
        'scaler_name': request_data.get('scalerName', 'None'),
        'support_filters': request_data.get('supportFilters', {}),
    }


async def trigger_temporal_workflow(workflow_input: dict) -> str:
    """Connect to Temporal and trigger the training workflow."""
    temporal_client = await client.Client.connect(
        target_host=TEMPORAL_HOST,
        namespace=TEMPORAL_NAMESPACE,
        tls=os.getenv('TEMPORAL_USE_TLS', False),
    )

    workflow_id = f'train-model-test-{uuid.uuid4()}'
    await temporal_client.execute_workflow(
        TEMPORAL_WORKFLOW,
        workflow_input,
        id=workflow_id,
        task_queue=TRAIN_TASK_QUEUE,
        execution_timeout=timedelta(minutes=5),
        run_timeout=timedelta(minutes=5),
        task_timeout=timedelta(minutes=5),
    )
    return workflow_id


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Run training workflow tests with different scenarios.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List available scenarios
  python scripts/run_training_test.py --list

  # Run a specific scenario
  python scripts/run_training_test.py --scenario linear-regression-basic

  # Run with a custom JSON file
  python scripts/run_training_test.py --scenario /path/to/custom-scenario.json

  # Run with a custom CSV data file
  python scripts/run_training_test.py --scenario linear-regression-basic --csv docs/other-data.csv
""",
    )
    parser.add_argument(
        '--scenario',
        '-s',
        type=str,
        help='Name of the test scenario (without .json) or path to a JSON file.',
    )
    parser.add_argument(
        '--csv',
        '-c',
        type=Path,
        default=DEFAULT_CSV_PATH,
        help=f'Path to the CSV data file (default: {DEFAULT_CSV_PATH}).',
    )
    parser.add_argument(
        '--list',
        '-l',
        action='store_true',
        help='List all available test scenarios and exit.',
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    # List scenarios and exit if requested
    if args.list:
        scenarios = list_available_scenarios()
        if scenarios:
            print('Available test scenarios:')
            for scenario in scenarios:
                print(f'  - {scenario}')
        else:
            print(f'No scenarios found in {TEST_SCENARIOS_DIR}')
        sys.exit(0)

    # Require scenario argument if not listing
    if not args.scenario:
        print('Error: --scenario is required. Use --list to see available scenarios.', file=sys.stderr)
        sys.exit(1)

    # Load scenario
    try:
        experiment_request = load_scenario(args.scenario)
        print(f"Loaded scenario: {args.scenario}")
    except FileNotFoundError as exc:
        print(str(exc), file=sys.stderr)
        sys.exit(1)

    # Upload CSV to MinIO
    try:
        uploaded_file_name = upload_to_minio(args.csv)
        print(f"Uploaded CSV to MinIO: {uploaded_file_name}")
    except subprocess.CalledProcessError as exc:
        print(f'Failed to upload file to MinIO: {exc}', file=sys.stderr)
        sys.exit(1)
    except FileNotFoundError as exc:
        print(str(exc), file=sys.stderr)
        sys.exit(1)

    # Insert experiment run
    try:
        experiment_run_id = insert_experiment_run(uploaded_file_name, experiment_request)
        print(f"Created experiment_run with ID: {experiment_run_id}")
    except psycopg2.Error as exc:
        print(f'Database error while inserting experiment_run: {exc}', file=sys.stderr)
        sys.exit(1)

    # Build and trigger workflow
    workflow_payload = build_workflow_payload(
        experiment_run_id=experiment_run_id,
        file_name=uploaded_file_name,
        request_data=experiment_request,
    )

    try:
        workflow_id = asyncio.run(trigger_temporal_workflow(workflow_payload))
    except Exception as exc:  # noqa: BLE001
        print(f'Failed to start Temporal workflow: {exc}', file=sys.stderr)
        sys.exit(1)

    print(
        json.dumps(
            {
                'scenario': args.scenario,
                'experiment_run_id': experiment_run_id,
                's3_object_name': uploaded_file_name,
                'workflow_id': workflow_id,
            },
            indent=2,
        )
    )


if __name__ == '__main__':
    main()