"""Tests for PingPulseOperator."""

import pytest
from unittest.mock import Mock, patch, MagicMock
from datetime import datetime

from pingpulse_airflow.operators.pingpulse import PingPulseOperator, PingPulseHeartbeatOperator


class TestPingPulseOperator:
    """Tests for PingPulseOperator class."""

    def test_init_required_params(self):
        """Test operator initialization with required parameters."""
        op = PingPulseOperator(
            task_id="test_ping",
            workflow_id="dwf123abc",
            stage_path="1",
        )
        assert op.workflow_id == "dwf123abc"
        assert op.stage_path == "1"
        assert op.start is False
        assert op.final is False
        assert op.workflow_type == "dynamic"
        assert op.include_context is True

    def test_init_all_params(self):
        """Test operator initialization with all parameters."""
        op = PingPulseOperator(
            task_id="test_ping",
            workflow_id="swf456def",
            stage_path="2.1",
            payload={"custom": "data"},
            start=True,
            final=False,
            workflow_type="static",
            api_key="ppk_test_key",
            base_url="https://custom.url",
            include_context=False,
        )
        assert op.workflow_id == "swf456def"
        assert op.stage_path == "2.1"
        assert op.payload == {"custom": "data"}
        assert op.start is True
        assert op.workflow_type == "static"
        assert op.api_key == "ppk_test_key"
        assert op.include_context is False

    def test_template_fields(self):
        """Test that template fields are properly defined."""
        assert "workflow_id" in PingPulseOperator.template_fields
        assert "stage_path" in PingPulseOperator.template_fields
        assert "payload" in PingPulseOperator.template_fields
        assert "api_key" in PingPulseOperator.template_fields

    def test_ui_color(self):
        """Test that UI color is set (PingPulse green)."""
        assert PingPulseOperator.ui_color == "#10B981"

    @patch("pingpulse_airflow.operators.pingpulse.PingPulseHook")
    def test_execute_sends_ping(self, MockHook):
        """Test that execute sends a ping via the hook."""
        mock_hook_instance = MagicMock()
        mock_hook_instance.send_ping.return_value = {"status": "success"}
        MockHook.return_value = mock_hook_instance

        op = PingPulseOperator(
            task_id="test_ping",
            workflow_id="dwf123abc",
            stage_path="1",
            api_key="ppk_test",
        )

        # Create mock context
        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "test_run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1),
        }

        result = op.execute(context)

        mock_hook_instance.send_ping.assert_called_once()
        call_kwargs = mock_hook_instance.send_ping.call_args[1]
        assert call_kwargs["workflow_id"] == "dwf123abc"
        assert call_kwargs["stage_path"] == "1"
        assert result == {"status": "success"}

    @patch("pingpulse_airflow.operators.pingpulse.PingPulseHook")
    def test_execute_includes_airflow_context(self, MockHook):
        """Test that execute includes Airflow context in payload."""
        mock_hook_instance = MagicMock()
        mock_hook_instance.send_ping.return_value = {"status": "success"}
        MockHook.return_value = mock_hook_instance

        op = PingPulseOperator(
            task_id="test_ping",
            workflow_id="dwf123abc",
            stage_path="1",
            api_key="ppk_test",
            include_context=True,
        )

        context = {
            "dag": MagicMock(dag_id="my_dag"),
            "task": MagicMock(task_id="my_task"),
            "run_id": "run_456",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=2),
        }

        op.execute(context)

        call_kwargs = mock_hook_instance.send_ping.call_args[1]
        payload = call_kwargs["payload"]
        assert "airflow" in payload
        assert payload["airflow"]["dag_id"] == "my_dag"
        assert payload["airflow"]["task_id"] == "my_task"
        assert payload["airflow"]["run_id"] == "run_456"

    @patch("pingpulse_airflow.operators.pingpulse.PingPulseHook")
    def test_execute_without_context(self, MockHook):
        """Test that execute works without including context."""
        mock_hook_instance = MagicMock()
        mock_hook_instance.send_ping.return_value = {"status": "success"}
        MockHook.return_value = mock_hook_instance

        op = PingPulseOperator(
            task_id="test_ping",
            workflow_id="dwf123abc",
            stage_path="1",
            api_key="ppk_test",
            include_context=False,
        )

        context = {
            "dag": MagicMock(dag_id="my_dag"),
            "task": MagicMock(task_id="my_task"),
            "run_id": "run_456",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1),
        }

        op.execute(context)

        call_kwargs = mock_hook_instance.send_ping.call_args[1]
        payload = call_kwargs["payload"]
        assert "airflow" not in payload


class TestPingPulseHeartbeatOperator:
    """Tests for PingPulseHeartbeatOperator class."""

    def test_init(self):
        """Test heartbeat operator initialization."""
        op = PingPulseHeartbeatOperator(
            task_id="heartbeat",
            monitor_id="hb123abc",
        )
        assert op.monitor_id == "hb123abc"
        assert op.payload == {}

    @patch("pingpulse_airflow.operators.pingpulse.PingPulseHook")
    def test_execute_sends_heartbeat(self, MockHook):
        """Test that execute sends heartbeat via hook."""
        mock_hook_instance = MagicMock()
        mock_hook_instance.send_heartbeat.return_value = {"status": "ok"}
        MockHook.return_value = mock_hook_instance

        op = PingPulseHeartbeatOperator(
            task_id="heartbeat",
            monitor_id="hb123abc",
            api_key="ppk_test",
        )

        context = {
            "dag": MagicMock(),
            "task": MagicMock(),
            "run_id": "test",
            "execution_date": datetime(2024, 1, 1),
            "ti": MagicMock(),
        }

        result = op.execute(context)

        mock_hook_instance.send_heartbeat.assert_called_once_with(
            monitor_id="hb123abc",
            payload={},
        )
        assert result == {"status": "ok"}
