"""Tests for PingPulse callbacks."""

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

from pingpulse_airflow.callbacks import (
    pingpulse_callback,
    pingpulse_success,
    pingpulse_failure,
    pingpulse_dag_callbacks,
)


class TestPingPulseCallback:
    """Tests for pingpulse_callback function."""

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_callback_sends_ping(self, MockHook):
        """Test that callback sends ping with correct parameters."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callback = pingpulse_callback(
            workflow_id="dwf123abc",
            stage_path="1",
            status="success",
            api_key="ppk_test",
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=10.5, log=MagicMock()),
        }

        callback(context)

        mock_hook.send_ping.assert_called_once()
        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["workflow_id"] == "dwf123abc"
        assert call_kwargs["stage_path"] == "1"
        assert call_kwargs["payload"]["status"] == "success"
        assert call_kwargs["payload"]["airflow"]["dag_id"] == "test_dag"

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_callback_includes_duration(self, MockHook):
        """Test that callback includes task duration."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callback = pingpulse_callback(
            workflow_id="dwf123abc",
            stage_path="1",
            api_key="ppk_test",
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=45.678, log=MagicMock()),
        }

        callback(context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["payload"]["duration"] == 45.678

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_callback_includes_exception_on_failure(self, MockHook):
        """Test that callback includes exception info on failure."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callback = pingpulse_callback(
            workflow_id="dwf123abc",
            stage_path="1",
            status="failed",
            api_key="ppk_test",
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=5.0, log=MagicMock()),
            "exception": ValueError("Test error message"),
        }

        callback(context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["payload"]["error"] == "Test error message"

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_callback_does_not_fail_on_hook_error(self, MockHook):
        """Test that callback doesn't raise even if hook fails."""
        mock_hook = MagicMock()
        mock_hook.send_ping.side_effect = Exception("Network error")
        MockHook.return_value = mock_hook

        callback = pingpulse_callback(
            workflow_id="dwf123abc",
            stage_path="1",
            api_key="ppk_test",
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=5.0, log=MagicMock()),
        }

        # Should not raise
        callback(context)


class TestPingPulseSuccess:
    """Tests for pingpulse_success function."""

    def test_returns_callable(self):
        """Test that pingpulse_success returns a callable."""
        callback = pingpulse_success("dwf123abc", "1")
        assert callable(callback)

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_sends_success_status(self, MockHook):
        """Test that callback sends success status."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callback = pingpulse_success(
            "dwf123abc", "1", api_key="ppk_test"
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=5.0, log=MagicMock()),
        }

        callback(context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["payload"]["status"] == "success"

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_respects_final_flag(self, MockHook):
        """Test that final flag is passed correctly."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callback = pingpulse_success(
            "dwf123abc", "3", final=True, api_key="ppk_test"
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=5.0, log=MagicMock()),
        }

        callback(context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["final"] is True


class TestPingPulseFailure:
    """Tests for pingpulse_failure function."""

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_sends_failed_status(self, MockHook):
        """Test that callback sends failed status."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callback = pingpulse_failure(
            "dwf123abc", "2", api_key="ppk_test"
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=5.0, log=MagicMock()),
            "exception": Exception("Task failed"),
        }

        callback(context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["payload"]["status"] == "failed"

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_never_sets_final_on_failure(self, MockHook):
        """Test that final is never True on failure."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        # Even if user passes final=True (shouldn't be allowed, but test anyway)
        callback = pingpulse_failure(
            "dwf123abc", "3", api_key="ppk_test"
        )

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "task": MagicMock(task_id="test_task"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
            "ti": MagicMock(try_number=1, duration=5.0, log=MagicMock()),
        }

        callback(context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["final"] is False


class TestPingPulseDagCallbacks:
    """Tests for pingpulse_dag_callbacks function."""

    def test_returns_dict_with_callbacks(self):
        """Test that function returns dict with both callbacks."""
        callbacks = pingpulse_dag_callbacks("dwf123abc")
        assert "on_success" in callbacks
        assert "on_failure" in callbacks
        assert callable(callbacks["on_success"])
        assert callable(callbacks["on_failure"])

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_dag_success_callback(self, MockHook):
        """Test DAG success callback sends correct ping."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callbacks = pingpulse_dag_callbacks("dwf123abc", api_key="ppk_test")

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
        }

        callbacks["on_success"](context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["payload"]["status"] == "dag_success"
        assert call_kwargs["final"] is True

    @patch("pingpulse_airflow.callbacks.PingPulseHook")
    def test_dag_failure_callback(self, MockHook):
        """Test DAG failure callback sends correct ping."""
        mock_hook = MagicMock()
        MockHook.return_value = mock_hook

        callbacks = pingpulse_dag_callbacks("dwf123abc", api_key="ppk_test")

        context = {
            "dag": MagicMock(dag_id="test_dag"),
            "run_id": "run_123",
            "execution_date": datetime(2024, 1, 15),
        }

        callbacks["on_failure"](context)

        call_kwargs = mock_hook.send_ping.call_args[1]
        assert call_kwargs["payload"]["status"] == "dag_failed"
        assert call_kwargs["final"] is False
