"""
PingPulse Task Decorator for Apache Airflow.

Provides a decorator to automatically ping PingPulse on task success/failure.
"""

from functools import wraps
from typing import Any, Callable, Dict, Optional
import time

from airflow.decorators import task
from airflow.exceptions import AirflowException

from pingpulse_airflow.hooks.pingpulse import PingPulseHook


def pingpulse_task(
    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,
    include_duration: bool = True,
    ping_on_failure: bool = True,
    **task_kwargs,
) -> Callable:
    """
    Decorator that wraps a Python function to auto-ping PingPulse.

    Automatically sends a ping on successful completion, and optionally
    on failure. Includes execution duration in the payload.

    Example:
        from pingpulse_airflow import pingpulse_task

        @pingpulse_task(
            workflow_id='dwf123abc',
            stage_path='1',
            start=True,
        )
        def extract_data():
            # Your code here
            return {'records': 1000}

        @pingpulse_task(
            workflow_id='dwf123abc',
            stage_path='2',
        )
        def transform_data(records):
            # Process records
            return transformed_data

        @pingpulse_task(
            workflow_id='dwf123abc',
            stage_path='3',
            final=True,
        )
        def load_data(data):
            # Load to destination
            return {'status': 'complete'}

    Args:
        workflow_id: PingPulse workflow ID
        stage_path: Stage path/order
        start: Set True for first stage (starts new instance)
        final: Set True for last stage (completes instance)
        workflow_type: 'dynamic' or 'static'
        pingpulse_conn_id: Airflow connection ID
        api_key: Direct API key (overrides connection)
        base_url: Direct base URL (overrides connection)
        include_duration: Include execution duration in payload
        ping_on_failure: Also ping on task failure (with status='failed')
        **task_kwargs: Additional kwargs passed to @task decorator

    Returns:
        Decorated function that auto-pings PingPulse
    """

    def decorator(func: Callable) -> Callable:
        @task(**task_kwargs)
        @wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            hook = PingPulseHook(
                pingpulse_conn_id=pingpulse_conn_id,
                api_key=api_key,
                base_url=base_url,
            )

            start_time = time.time()
            result = None
            error = None

            try:
                result = func(*args, **kwargs)
                return result

            except Exception as e:
                error = e
                raise

            finally:
                duration = time.time() - start_time

                # Build payload
                payload: Dict[str, Any] = {
                    "task_name": func.__name__,
                }

                if include_duration:
                    payload["duration"] = round(duration, 3)

                if error:
                    if ping_on_failure:
                        payload["status"] = "failed"
                        payload["error"] = str(error)[:500]  # Truncate long errors
                        try:
                            hook.send_ping(
                                workflow_id=workflow_id,
                                stage_path=stage_path,
                                payload=payload,
                                start=start,
                                final=False,  # Don't mark final on failure
                                workflow_type=workflow_type,
                            )
                        except Exception:
                            pass  # Don't mask original error
                else:
                    payload["status"] = "success"
                    if result is not None and isinstance(result, dict):
                        # Include return value if it's a dict
                        payload["result"] = result

                    hook.send_ping(
                        workflow_id=workflow_id,
                        stage_path=stage_path,
                        payload=payload,
                        start=start,
                        final=final,
                        workflow_type=workflow_type,
                    )

        return wrapper

    return decorator


def pingpulse_python_callable(
    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,
    include_duration: bool = True,
    ping_on_failure: bool = True,
) -> Callable:
    """
    Decorator for use with PythonOperator's python_callable.

    Unlike pingpulse_task, this doesn't use @task decorator internally,
    making it suitable for traditional PythonOperator usage.

    Example:
        from pingpulse_airflow.decorators import pingpulse_python_callable

        @pingpulse_python_callable(
            workflow_id='dwf123abc',
            stage_path='1',
        )
        def my_function(**context):
            return {'result': 'data'}

        task = PythonOperator(
            task_id='my_task',
            python_callable=my_function,
        )
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            hook = PingPulseHook(
                pingpulse_conn_id=pingpulse_conn_id,
                api_key=api_key,
                base_url=base_url,
            )

            start_time = time.time()
            result = None
            error = None

            try:
                result = func(*args, **kwargs)
                return result

            except Exception as e:
                error = e
                raise

            finally:
                duration = time.time() - start_time

                payload: Dict[str, Any] = {
                    "task_name": func.__name__,
                }

                if include_duration:
                    payload["duration"] = round(duration, 3)

                # Extract context if available
                context = kwargs.get("context") or (
                    kwargs if "ti" in kwargs or "dag" in kwargs else None
                )
                if context:
                    if "dag" in context:
                        payload.setdefault("airflow", {})["dag_id"] = context["dag"].dag_id
                    if "task" in context:
                        payload.setdefault("airflow", {})["task_id"] = context["task"].task_id
                    if "run_id" in context:
                        payload.setdefault("airflow", {})["run_id"] = context["run_id"]

                if error:
                    if ping_on_failure:
                        payload["status"] = "failed"
                        payload["error"] = str(error)[:500]
                        try:
                            hook.send_ping(
                                workflow_id=workflow_id,
                                stage_path=stage_path,
                                payload=payload,
                                start=start,
                                final=False,
                                workflow_type=workflow_type,
                            )
                        except Exception:
                            pass
                else:
                    payload["status"] = "success"
                    if result is not None and isinstance(result, dict):
                        payload["result"] = result

                    hook.send_ping(
                        workflow_id=workflow_id,
                        stage_path=stage_path,
                        payload=payload,
                        start=start,
                        final=final,
                        workflow_type=workflow_type,
                    )

        return wrapper

    return decorator
