"""
Unit tests for Demand Response Service
Run with: pytest test_demand_response.py -v
"""

import pytest
import asyncio
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock, patch
import json

from demand_response_service import DemandResponseService


# Test fixtures
@pytest.fixture
def mock_db():
    """Mock MongoDB database"""
    db = MagicMock()

    # Mock collections
    db.demand_response_invitations = MagicMock()
    db.demand_response_events = MagicMock()
    db.demand_response_responses = MagicMock()
    db.auto_response_config = MagicMock()
    db.device_instructions = MagicMock()
    db.flexibility_snapshots = MagicMock()

    return db


@pytest.fixture
def mock_redis():
    """Mock Redis client"""
    redis = AsyncMock()
    redis.get = AsyncMock(return_value=None)
    redis.setex = AsyncMock()
    redis.delete = AsyncMock()
    redis.publish = AsyncMock()
    return redis


@pytest.fixture
def dr_service(mock_db, mock_redis):
    """Create DemandResponseService instance with mocks"""
    return DemandResponseService(mock_db, mock_redis)


# Test: Invitation Management

@pytest.mark.asyncio
async def test_send_invitation_with_auto_accept(dr_service, mock_db, mock_redis):
    """Test sending invitation with auto-accept enabled"""
    # Mock auto-response config (enabled)
    mock_db.auto_response_config.find_one = AsyncMock(return_value={
        "config_id": "default",
        "enabled": True
    })

    mock_db.demand_response_invitations.insert_one = AsyncMock()

    event_time = datetime.utcnow() + timedelta(hours=2)
    result = await dr_service.send_invitation(
        event_time=event_time,
        load_kwh=5.0,
        load_percentage=15.0,
        iots=["sensor_1", "sensor_2"],
        duration_minutes=59
    )

    assert "event_id" in result
    assert result["response"] == "YES"
    assert result["message"] == "Invitation created successfully"

    # Verify MongoDB insert was called
    mock_db.demand_response_invitations.insert_one.assert_called_once()

    # Verify Redis caching
    mock_redis.setex.assert_called()
    mock_redis.publish.assert_called()


@pytest.mark.asyncio
async def test_send_invitation_manual(dr_service, mock_db, mock_redis):
    """Test sending invitation with auto-accept disabled (manual mode)"""
    # Mock auto-response config (disabled)
    mock_db.auto_response_config.find_one = AsyncMock(return_value={
        "config_id": "default",
        "enabled": False
    })

    mock_db.demand_response_invitations.insert_one = AsyncMock()

    event_time = datetime.utcnow() + timedelta(hours=2)
    result = await dr_service.send_invitation(
        event_time=event_time,
        load_kwh=5.0,
        load_percentage=15.0,
        iots=["sensor_1", "sensor_2"],
        duration_minutes=59
    )

    assert result["response"] == "WAITING"


@pytest.mark.asyncio
async def test_answer_invitation_success(dr_service, mock_db, mock_redis):
    """Test answering an invitation successfully"""
    event_id = "test-event-123"

    # Mock get_invitation to return a valid invitation
    dr_service.get_invitation = AsyncMock(return_value={
        "event_id": event_id,
        "iots": ["sensor_1", "sensor_2"]
    })

    # Mock that device hasn't responded yet
    mock_db.demand_response_responses.find_one = AsyncMock(return_value=None)
    mock_db.demand_response_responses.insert_one = AsyncMock()
    mock_db.demand_response_responses.count_documents = AsyncMock(return_value=1)

    result = await dr_service.answer_invitation(
        event_id=event_id,
        iot_id="sensor_1",
        response="YES",
        committed_reduction_kw=2.5
    )

    assert result["success"] is True
    assert result["message"] == "Response recorded successfully"

    # Verify response was stored
    mock_db.demand_response_responses.insert_one.assert_called_once()
    mock_redis.delete.assert_called()
    mock_redis.publish.assert_called()


@pytest.mark.asyncio
async def test_answer_invitation_device_not_in_list(dr_service, mock_db, mock_redis):
    """Test answering invitation for device not in invitation list"""
    event_id = "test-event-123"

    dr_service.get_invitation = AsyncMock(return_value={
        "event_id": event_id,
        "iots": ["sensor_1", "sensor_2"]
    })

    result = await dr_service.answer_invitation(
        event_id=event_id,
        iot_id="sensor_3",  # Not in list
        response="YES"
    )

    assert result["success"] is False
    assert "not in invitation" in result["message"]


# Test: Event Execution

@pytest.mark.asyncio
async def test_schedule_event(dr_service, mock_db, mock_redis):
    """Test scheduling a DR event"""
    mock_db.demand_response_events.insert_one = AsyncMock()

    event_time = datetime.utcnow() + timedelta(hours=1)
    result = await dr_service.schedule_event(
        event_time=event_time,
        iots=["sensor_1", "sensor_2"],
        load_reduction_kw=5.0,
        duration_minutes=59
    )

    assert "event_id" in result
    assert result["message"] == "Event scheduled successfully"

    mock_db.demand_response_events.insert_one.assert_called_once()
    mock_redis.publish.assert_called()


@pytest.mark.asyncio
async def test_execute_event(dr_service, mock_db, mock_redis):
    """Test executing a DR event (spawns background task)"""
    event_id = "test-event-456"

    # Mock event document
    event = {
        "event_id": event_id,
        "start_time": datetime.utcnow(),
        "end_time": datetime.utcnow() + timedelta(minutes=59),
        "participating_devices": ["sensor_1"],
        "target_reduction_kw": 5.0
    }

    mock_db.demand_response_events.find_one = AsyncMock(return_value=event)
    mock_db.demand_response_events.update_one = AsyncMock()

    # Execute event (starts background task)
    await dr_service.execute_event(event_id)

    # Verify event status updated to active
    mock_db.demand_response_events.update_one.assert_called()
    mock_redis.publish.assert_called()

    # Verify task was created and stored
    assert event_id in dr_service.active_events

    # Cancel the task to prevent it from running
    task = dr_service.active_events[event_id]
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass


@pytest.mark.asyncio
async def test_cancel_event(dr_service, mock_db, mock_redis):
    """Test cancelling a running DR event"""
    event_id = "test-event-789"

    # Create a mock task
    mock_task = AsyncMock()
    mock_task.done = MagicMock(return_value=False)
    mock_task.cancel = MagicMock()

    dr_service.active_events[event_id] = mock_task

    # Mock database operations
    mock_db.demand_response_events.find_one = AsyncMock(return_value={
        "event_id": event_id,
        "status": "active"
    })
    mock_db.demand_response_events.update_one = AsyncMock()

    await dr_service.cancel_event(event_id)

    # Verify task was cancelled
    mock_task.cancel.assert_called_once()

    # Verify database updated
    mock_db.demand_response_events.update_one.assert_called()
    mock_redis.delete.assert_called()
    mock_redis.publish.assert_called()


# Test: Device Power Integration

@pytest.mark.asyncio
async def test_update_device_power_cache(dr_service):
    """Test updating device power cache"""
    dr_service.update_device_power_cache("sensor_1", 2.5)

    assert dr_service.device_power_cache["sensor_1"] == 2.5

    dr_service.update_device_power_cache("sensor_1", 3.0)
    assert dr_service.device_power_cache["sensor_1"] == 3.0


@pytest.mark.asyncio
async def test_get_device_power(dr_service):
    """Test getting device power from cache"""
    dr_service.device_power_cache["sensor_1"] = 2.5

    power = await dr_service.get_device_power("sensor_1")
    assert power == 2.5

    # Test non-existent device returns 0
    power = await dr_service.get_device_power("sensor_999")
    assert power == 0.0


# Test: Auto-Response Configuration

@pytest.mark.asyncio
async def test_get_auto_response_config_exists(dr_service, mock_db):
    """Test getting existing auto-response config"""
    mock_config = {
        "config_id": "default",
        "enabled": True,
        "max_reduction_percentage": 20.0
    }

    mock_db.auto_response_config.find_one = AsyncMock(return_value=mock_config)

    config = await dr_service.get_auto_response_config()

    assert config["enabled"] is True
    assert config["max_reduction_percentage"] == 20.0


@pytest.mark.asyncio
async def test_get_auto_response_config_creates_default(dr_service, mock_db):
    """Test creating default config when none exists"""
    mock_db.auto_response_config.find_one = AsyncMock(return_value=None)
    mock_db.auto_response_config.insert_one = AsyncMock()

    config = await dr_service.get_auto_response_config()

    assert config["enabled"] is False
    mock_db.auto_response_config.insert_one.assert_called_once()


@pytest.mark.asyncio
async def test_set_auto_response_config(dr_service, mock_db, mock_redis):
    """Test updating auto-response configuration"""
    mock_db.auto_response_config.update_one = AsyncMock()
    mock_db.auto_response_config.find_one = AsyncMock(return_value={
        "config_id": "default",
        "enabled": True,
        "max_reduction_percentage": 25.0
    })

    config = await dr_service.set_auto_response_config(
        enabled=True,
        max_reduction_percentage=25.0
    )

    assert config["enabled"] is True
    assert config["max_reduction_percentage"] == 25.0

    mock_db.auto_response_config.update_one.assert_called_once()
    mock_redis.delete.assert_called()


# Test: Auto-Response Processing

@pytest.mark.asyncio
async def test_process_auto_responses_disabled(dr_service, mock_db):
    """Test auto-response processing when disabled"""
    mock_db.auto_response_config.find_one = AsyncMock(return_value={
        "config_id": "default",
        "enabled": False
    })

    # Should return early without processing
    await dr_service.process_auto_responses()

    # No invitations should be queried
    mock_db.demand_response_invitations.find.assert_not_called()


@pytest.mark.asyncio
async def test_process_auto_responses_enabled(dr_service, mock_db, mock_redis):
    """Test auto-response processing when enabled"""
    # Mock enabled config
    mock_db.auto_response_config.find_one = AsyncMock(return_value={
        "config_id": "default",
        "enabled": True,
        "max_reduction_percentage": 20.0,
        "min_notice_minutes": 60
    })

    # Mock pending invitation
    future_time = datetime.utcnow() + timedelta(hours=2)
    mock_invitation = {
        "event_id": "test-event-auto",
        "event_time": future_time,
        "iots": ["sensor_1"]
    }

    dr_service.get_unanswered_invitations = AsyncMock(return_value=[mock_invitation])
    dr_service.get_device_power = AsyncMock(return_value=5.0)
    dr_service.answer_invitation = AsyncMock(return_value={"success": True})

    mock_db.demand_response_responses.find_one = AsyncMock(return_value=None)

    await dr_service.process_auto_responses()

    # Should have auto-responded
    dr_service.answer_invitation.assert_called_once()


# Test: Flexibility Calculation

@pytest.mark.asyncio
async def test_get_current_flexibility(dr_service, mock_db, mock_redis):
    """Test calculating current flexibility"""
    # Mock device with instructions
    mock_device = {
        "device_id": "sensor_1",
        "instructions": {
            str(datetime.utcnow().hour): "participation"
        }
    }

    async def mock_cursor():
        yield mock_device

    mock_db.device_instructions.find = MagicMock(return_value=mock_cursor())
    mock_db.flexibility_snapshots.insert_one = AsyncMock()

    # Set device power in cache
    dr_service.device_power_cache["sensor_1"] = 5.0

    result = await dr_service.get_current_flexibility()

    assert result["total_flexibility_kw"] == 5.0
    assert len(result["devices"]) == 1
    assert result["devices"][0]["device_id"] == "sensor_1"

    mock_db.flexibility_snapshots.insert_one.assert_called_once()
    mock_redis.setex.assert_called()


# Test: Device Instructions

@pytest.mark.asyncio
async def test_update_device_instructions(dr_service, mock_db):
    """Test updating device DR instructions"""
    mock_db.device_instructions.update_one = AsyncMock()

    instructions = {
        "0": "participation",
        "1": "shifting",
        "2": "off"
    }

    await dr_service.update_device_instructions("sensor_1", instructions)

    mock_db.device_instructions.update_one.assert_called_once()


@pytest.mark.asyncio
async def test_get_device_instructions_single(dr_service, mock_db):
    """Test getting instructions for single device"""
    mock_instructions = {
        "device_id": "sensor_1",
        "instructions": {"0": "participation"}
    }

    mock_db.device_instructions.find_one = AsyncMock(return_value=mock_instructions)

    result = await dr_service.get_device_instructions("sensor_1")

    assert result["device_id"] == "sensor_1"
    assert "instructions" in result


# Test: Analytics

@pytest.mark.asyncio
async def test_get_performance_analytics(dr_service, mock_db):
    """Test getting performance analytics"""
    # Mock completed events
    mock_events = [
        {"actual_reduction_kw": 5.0, "target_reduction_kw": 6.0},
        {"actual_reduction_kw": 4.5, "target_reduction_kw": 5.0}
    ]

    mock_cursor = AsyncMock()
    mock_cursor.to_list = AsyncMock(return_value=mock_events)

    mock_db.demand_response_events.find = MagicMock(return_value=mock_cursor)

    analytics = await dr_service.get_performance_analytics(days=30)

    assert analytics["total_events"] == 2
    assert analytics["total_reduction_kwh"] == 9.5
    assert analytics["total_target_kwh"] == 11.0
    assert analytics["achievement_rate"] > 0


@pytest.mark.asyncio
async def test_get_performance_analytics_no_events(dr_service, mock_db):
    """Test analytics with no completed events"""
    mock_cursor = AsyncMock()
    mock_cursor.to_list = AsyncMock(return_value=[])

    mock_db.demand_response_events.find = MagicMock(return_value=mock_cursor)

    analytics = await dr_service.get_performance_analytics(days=30)

    assert analytics["total_events"] == 0
    assert analytics["total_reduction_kwh"] == 0.0
    assert analytics["achievement_rate"] == 0.0


# Integration-style tests

@pytest.mark.asyncio
async def test_full_invitation_workflow(dr_service, mock_db, mock_redis):
    """Test complete invitation workflow from creation to response"""
    # Step 1: Create invitation
    mock_db.auto_response_config.find_one = AsyncMock(return_value={
        "config_id": "default",
        "enabled": False
    })
    mock_db.demand_response_invitations.insert_one = AsyncMock()

    event_time = datetime.utcnow() + timedelta(hours=2)
    invite_result = await dr_service.send_invitation(
        event_time=event_time,
        load_kwh=5.0,
        load_percentage=15.0,
        iots=["sensor_1", "sensor_2"],
        duration_minutes=59
    )

    event_id = invite_result["event_id"]
    assert invite_result["response"] == "WAITING"

    # Step 2: Answer invitation for device 1
    dr_service.get_invitation = AsyncMock(return_value={
        "event_id": event_id,
        "iots": ["sensor_1", "sensor_2"]
    })
    mock_db.demand_response_responses.find_one = AsyncMock(return_value=None)
    mock_db.demand_response_responses.insert_one = AsyncMock()
    mock_db.demand_response_responses.count_documents = AsyncMock(side_effect=[1, 1, 2, 2])
    mock_db.demand_response_invitations.update_one = AsyncMock()

    answer1 = await dr_service.answer_invitation(event_id, "sensor_1", "YES", 2.5)
    assert answer1["success"] is True

    # Step 3: Answer invitation for device 2
    answer2 = await dr_service.answer_invitation(event_id, "sensor_2", "YES", 2.5)
    assert answer2["success"] is True

    # Verify final invitation update was called (all devices responded)
    assert mock_db.demand_response_invitations.update_one.call_count >= 1


if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])