"""
PingPulse Operator for Apache Airflow.

Provides an operator to send pings to PingPulse as explicit DAG tasks.
"""

from typing import Any, Dict, Optional, Sequence
from airflow.models import BaseOperator
from airflow.utils.context import Context

from pingpulse_airflow.hooks.pingpulse import PingPulseHook


class PingPulseOperator(BaseOperator):
    """
    Send a ping to PingPulse workflow monitoring.

    Use this operator to create explicit ping tasks in your DAG.
    The ping will be visible as a separate task node.

    Example:
        from pingpulse_airflow import PingPulseOperator

        # Basic usage with Airflow connection
        ping_start = PingPulseOperator(
            task_id='ping_workflow_start',
            workflow_id='dwf123abc',
            stage_path='1',
            start=True,
        )

        # With direct API key
        ping_complete = PingPulseOperator(
            task_id='ping_workflow_complete',
            workflow_id='dwf123abc',
            stage_path='3',
            final=True,
            api_key='ppk_xxx_xxx',
        )

        # With custom payload
        ping_metrics = PingPulseOperator(
            task_id='ping_with_metrics',
            workflow_id='dwf123abc',
            stage_path='2',
            payload={
                'duration': '{{ ti.duration }}',
                'records_processed': 1000,
            },
        )

        extract >> ping_start >> transform >> ping_metrics >> load >> ping_complete
    """

    # Template fields for Jinja rendering
    template_fields: Sequence[str] = (
        "workflow_id",
        "stage_path",
        "payload",
        "api_key",
    )

    # UI color (PingPulse green)
    ui_color = "#10B981"
    ui_fgcolor = "#FFFFFF"

    def __init__(
        self,
        *,
        workflow_id: str,
        stage_path: str,
        payload: Optional[Dict[str, Any]] = None,
        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_context: bool = True,
        **kwargs,
    ) -> None:
        """
        Initialize PingPulse operator.

        Args:
            workflow_id: PingPulse workflow ID (e.g., 'dwf123abc', 'swf456def')
            stage_path: Stage path/order (e.g., '1', '2', '2.1' for branches)
            payload: Optional dict of metrics/metadata to send
            start: Set True to start a new workflow instance
            final: Set True to mark workflow instance as complete
            workflow_type: 'dynamic' (dhk endpoint) or 'static' (shk endpoint)
            pingpulse_conn_id: Airflow connection ID for PingPulse
            api_key: Direct API key (overrides connection)
            base_url: Direct base URL (overrides connection)
            include_context: Include Airflow context (dag_id, task_id, run_id) in payload
        """
        super().__init__(**kwargs)
        self.workflow_id = workflow_id
        self.stage_path = stage_path
        self.payload = payload or {}
        self.start = start
        self.final = final
        self.workflow_type = workflow_type
        self.pingpulse_conn_id = pingpulse_conn_id
        self.api_key = api_key
        self.base_url = base_url
        self.include_context = include_context

    def execute(self, context: Context) -> Dict[str, Any]:
        """
        Execute the PingPulse ping.

        Args:
            context: Airflow task context

        Returns:
            Response from PingPulse API
        """
        hook = PingPulseHook(
            pingpulse_conn_id=self.pingpulse_conn_id,
            api_key=self.api_key,
            base_url=self.base_url,
        )

        # Build payload with optional Airflow context
        payload = dict(self.payload)

        if self.include_context:
            payload.setdefault("airflow", {})
            payload["airflow"].update({
                "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,
            })

        self.log.info(
            f"Sending PingPulse ping: workflow={self.workflow_id}, "
            f"stage={self.stage_path}, start={self.start}, final={self.final}"
        )

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

        self.log.info(f"PingPulse response: {response}")
        return response


class PingPulseHeartbeatOperator(BaseOperator):
    """
    Send a heartbeat ping to PingPulse.

    Use this for cron job / scheduled task monitoring.

    Example:
        from pingpulse_airflow.operators.pingpulse import PingPulseHeartbeatOperator

        heartbeat = PingPulseHeartbeatOperator(
            task_id='send_heartbeat',
            monitor_id='hb123abc',
        )
    """

    template_fields: Sequence[str] = ("monitor_id", "payload", "api_key")
    ui_color = "#10B981"
    ui_fgcolor = "#FFFFFF"

    def __init__(
        self,
        *,
        monitor_id: str,
        payload: Optional[Dict[str, Any]] = None,
        pingpulse_conn_id: str = "pingpulse_default",
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        **kwargs,
    ) -> None:
        """
        Initialize heartbeat operator.

        Args:
            monitor_id: PingPulse heartbeat monitor ID (e.g., 'hb123abc')
            payload: Optional dict of metadata
            pingpulse_conn_id: Airflow connection ID
            api_key: Direct API key (overrides connection)
            base_url: Direct base URL (overrides connection)
        """
        super().__init__(**kwargs)
        self.monitor_id = monitor_id
        self.payload = payload or {}
        self.pingpulse_conn_id = pingpulse_conn_id
        self.api_key = api_key
        self.base_url = base_url

    def execute(self, context: Context) -> Dict[str, Any]:
        """Execute the heartbeat ping."""
        hook = PingPulseHook(
            pingpulse_conn_id=self.pingpulse_conn_id,
            api_key=self.api_key,
            base_url=self.base_url,
        )

        self.log.info(f"Sending PingPulse heartbeat: monitor={self.monitor_id}")

        response = hook.send_heartbeat(
            monitor_id=self.monitor_id,
            payload=self.payload,
        )

        self.log.info(f"PingPulse heartbeat response: {response}")
        return response
