"""
PingPulse Callbacks for Apache Airflow.

Provides callback functions for on_success_callback and on_failure_callback.
Use these to add PingPulse monitoring to existing tasks without modification.
"""

from typing import Any, Callable, Dict, Optional

from pingpulse_airflow.hooks.pingpulse import PingPulseHook


def pingpulse_callback(
    workflow_id: str,
    stage_path: str,
    *,
    status: str = "success",
    start: bool = False,
    final: bool = False,
    workflow_type: str = "dynamic",
    pingpulse_conn_id: str = "pingpulse_default",
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    extra_payload: Optional[Dict[str, Any]] = None,
) -> Callable:
    """
    Generic callback factory for PingPulse pings.

    Example:
        from pingpulse_airflow import pingpulse_callback

        task = PythonOperator(
            task_id='my_task',
            python_callable=my_func,
            on_success_callback=pingpulse_callback(
                'dwf123abc', '1', status='success'
            ),
            on_failure_callback=pingpulse_callback(
                'dwf123abc', '1', status='failed'
            ),
        )

    Args:
        workflow_id: PingPulse workflow ID
        stage_path: Stage path/order
        status: Status to send ('success', 'failed', etc.)
        start: Set True for first stage
        final: Set True for last stage
        workflow_type: 'dynamic' or 'static'
        pingpulse_conn_id: Airflow connection ID
        api_key: Direct API key
        base_url: Direct base URL
        extra_payload: Additional fields to include in payload

    Returns:
        Callback function compatible with Airflow callbacks
    """

    def callback(context: Dict[str, Any]) -> None:
        hook = PingPulseHook(
            pingpulse_conn_id=pingpulse_conn_id,
            api_key=api_key,
            base_url=base_url,
        )

        # Build payload from Airflow context
        payload: Dict[str, Any] = {
            "status": status,
            "airflow": {
                "dag_id": context["dag"].dag_id,
                "task_id": context["task"].task_id,
                "run_id": context["run_id"],
                "execution_date": str(context["execution_date"]),
                "try_number": context["ti"].try_number,
            },
        }

        # Include task duration if available
        ti = context.get("ti")
        if ti and hasattr(ti, "duration") and ti.duration:
            payload["duration"] = round(ti.duration, 3)

        # Include exception info for failures
        exception = context.get("exception")
        if exception:
            payload["error"] = str(exception)[:500]

        # Merge extra payload
        if extra_payload:
            payload.update(extra_payload)

        try:
            hook.send_ping(
                workflow_id=workflow_id,
                stage_path=stage_path,
                payload=payload,
                start=start,
                final=final if status == "success" else False,  # Don't finalize on failure
                workflow_type=workflow_type,
            )
        except Exception as e:
            # Log but don't fail the task
            context["ti"].log.warning(f"PingPulse callback failed: {e}")

    return callback


def pingpulse_success(
    workflow_id: str,
    stage_path: str,
    *,
    start: bool = False,
    final: bool = False,
    workflow_type: str = "dynamic",
    pingpulse_conn_id: str = "pingpulse_default",
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    extra_payload: Optional[Dict[str, Any]] = None,
) -> Callable:
    """
    Create an on_success_callback that pings PingPulse.

    Example:
        from pingpulse_airflow import pingpulse_success

        # Basic usage
        task = PythonOperator(
            task_id='extract',
            python_callable=extract_data,
            on_success_callback=pingpulse_success('dwf123abc', '1', start=True),
        )

        # Final stage
        load_task = PythonOperator(
            task_id='load',
            python_callable=load_data,
            on_success_callback=pingpulse_success('dwf123abc', '3', final=True),
        )

    Args:
        workflow_id: PingPulse workflow ID
        stage_path: Stage path/order
        start: Set True for first stage
        final: Set True for last stage
        workflow_type: 'dynamic' or 'static'
        pingpulse_conn_id: Airflow connection ID
        api_key: Direct API key
        base_url: Direct base URL
        extra_payload: Additional fields to include

    Returns:
        Callback function for on_success_callback
    """
    return pingpulse_callback(
        workflow_id=workflow_id,
        stage_path=stage_path,
        status="success",
        start=start,
        final=final,
        workflow_type=workflow_type,
        pingpulse_conn_id=pingpulse_conn_id,
        api_key=api_key,
        base_url=base_url,
        extra_payload=extra_payload,
    )


def pingpulse_failure(
    workflow_id: str,
    stage_path: str,
    *,
    start: bool = False,
    workflow_type: str = "dynamic",
    pingpulse_conn_id: str = "pingpulse_default",
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    extra_payload: Optional[Dict[str, Any]] = None,
) -> Callable:
    """
    Create an on_failure_callback that pings PingPulse.

    Example:
        from pingpulse_airflow import pingpulse_success, pingpulse_failure

        task = PythonOperator(
            task_id='transform',
            python_callable=transform_data,
            on_success_callback=pingpulse_success('dwf123abc', '2'),
            on_failure_callback=pingpulse_failure('dwf123abc', '2'),
        )

    Args:
        workflow_id: PingPulse workflow ID
        stage_path: Stage path/order
        start: Set True for first stage
        workflow_type: 'dynamic' or 'static'
        pingpulse_conn_id: Airflow connection ID
        api_key: Direct API key
        base_url: Direct base URL
        extra_payload: Additional fields to include

    Returns:
        Callback function for on_failure_callback
    """
    return pingpulse_callback(
        workflow_id=workflow_id,
        stage_path=stage_path,
        status="failed",
        start=start,
        final=False,  # Never finalize on failure
        workflow_type=workflow_type,
        pingpulse_conn_id=pingpulse_conn_id,
        api_key=api_key,
        base_url=base_url,
        extra_payload=extra_payload,
    )


def pingpulse_dag_callbacks(
    workflow_id: str,
    *,
    workflow_type: str = "dynamic",
    pingpulse_conn_id: str = "pingpulse_default",
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
) -> Dict[str, Callable]:
    """
    Create DAG-level callbacks for workflow start/complete.

    Example:
        from pingpulse_airflow.callbacks import pingpulse_dag_callbacks

        callbacks = pingpulse_dag_callbacks('dwf123abc')

        with DAG(
            'my_dag',
            on_success_callback=callbacks['on_success'],
            on_failure_callback=callbacks['on_failure'],
            ...
        ) as dag:
            ...

    Args:
        workflow_id: PingPulse workflow ID
        workflow_type: 'dynamic' or 'static'
        pingpulse_conn_id: Airflow connection ID
        api_key: Direct API key
        base_url: Direct base URL

    Returns:
        Dict with 'on_success' and 'on_failure' callbacks
    """

    def on_dag_success(context: Dict[str, Any]) -> None:
        hook = PingPulseHook(
            pingpulse_conn_id=pingpulse_conn_id,
            api_key=api_key,
            base_url=base_url,
        )

        payload = {
            "status": "dag_success",
            "airflow": {
                "dag_id": context["dag"].dag_id,
                "run_id": context["run_id"],
                "execution_date": str(context["execution_date"]),
            },
        }

        try:
            hook.send_ping(
                workflow_id=workflow_id,
                stage_path="dag_complete",
                payload=payload,
                final=True,
                workflow_type=workflow_type,
            )
        except Exception as e:
            print(f"PingPulse DAG success callback failed: {e}")

    def on_dag_failure(context: Dict[str, Any]) -> None:
        hook = PingPulseHook(
            pingpulse_conn_id=pingpulse_conn_id,
            api_key=api_key,
            base_url=base_url,
        )

        payload = {
            "status": "dag_failed",
            "airflow": {
                "dag_id": context["dag"].dag_id,
                "run_id": context["run_id"],
                "execution_date": str(context["execution_date"]),
            },
        }

        try:
            hook.send_ping(
                workflow_id=workflow_id,
                stage_path="dag_failed",
                payload=payload,
                final=False,
                workflow_type=workflow_type,
            )
        except Exception as e:
            print(f"PingPulse DAG failure callback failed: {e}")

    return {
        "on_success": on_dag_success,
        "on_failure": on_dag_failure,
    }
