import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from src.aeros_equipment.service import update_equipment_for_simulation
from src.aeros_equipment.model import AerosEquipment, AerosEquipmentGroup, ReliabilityPredictNonRepairable
from src.aeros_simulation.model import AerosNode
import logging

@pytest.mark.asyncio
async def test_update_equipment_for_simulation_basic(mocker):
    # Mock parameters
    db_session = AsyncMock()
    aeros_db_session = AsyncMock()
    project_name = "TestProject"
    simulation_duration = 8760
    overhaul_duration = 72
    overhaul_interval = 8000
    offset = 0
    schematic_name = "TestSchematic"
    
    # Mock get_aeros_schematic_by_name
    mock_schematic = MagicMock()
    mock_schematic.id = 1
    mocker.patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic)
    
    # Mock db_session.execute for equipments
    mock_eq = MagicMock()
    mock_eq.location_tag = "EQ001"
    
    mock_eq_result = MagicMock()
    mock_eq_result.scalars.return_value.all.return_value = [mock_eq]
    
    # Mock db_session.execute for rbd_group
    mock_group = MagicMock()
    mock_group.group_name = "GROUP001"
    
    mock_group_result = MagicMock()
    mock_group_result.scalars.return_value.all.return_value = [mock_group]
    
    # Mock db_session.execute for non_repairable
    mock_nr_result = MagicMock()
    mock_nr_result.scalars.return_value.all.return_value = []

    db_session.execute.side_effect = [
        mock_eq_result,    # for equipments
        mock_group_result, # for rbd_group
        mock_nr_result     # for non_repairable
    ]
    
    # Mock get_aeros_equipment_by_location_tag
    mock_nodes_data = [
        {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0},
        {"equipmentName": "GROUP001", "relDisType": "Fixed", "relDisP1": 0}
    ]
    mocker.patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data)
    
    # Mock get_asset_batch
    mock_reliability_data = {
        "EQ001": {
            "relDisType": "Weibull2",
            "relDisP1": 1.5,
            "relDisP2": 10000,
            "cmDisType": "Fixed",
            "cmDisP1": 5,
            "cmDisP2": 0
        }
    }
    mocker.patch("src.aeros_equipment.service.get_asset_batch", return_value=mock_reliability_data)
    
    # Mock update_oh_interval_offset
    mocker.patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None)
    
    # Mock update_node
    mock_update_node = mocker.patch("src.aeros_equipment.service.update_node", return_value=None)
    
    # Run the function
    results = await update_equipment_for_simulation(
        db_session=db_session,
        aeros_db_session=aeros_db_session,
        project_name=project_name,
        simulation_duration=simulation_duration,
        overhaul_duration=overhaul_duration,
        overhaul_interval=overhaul_interval,
        offset=offset,
        schematic_name=schematic_name
    )
    
    # Assertions
    assert "EQ001" in results
    assert results["EQ001"]["distribution"] == "Weibull2"
    assert results["EQ001"]["oh_duration"] == overhaul_duration
    
    # Check if update_node was called with correct arguments
    called_args = mock_update_node.call_args[1]
    equipment_nodes = called_args["equipment_nodes"]
    
    # Find EQ001 in equipment_nodes
    eq001_node = next(node for node in equipment_nodes if node["equipmentName"] == "EQ001")
    assert eq001_node["relDisType"] == "Weibull2"
    assert eq001_node["ohDisP1"] == overhaul_duration
    assert eq001_node["relDisUnitCode"] == "UHour"
    assert eq001_node["ohDisUnitCode"] == "UHour"

@pytest.mark.asyncio
async def test_update_equipment_for_simulation_custom_input(mocker):
    # Mock parameters
    db_session = AsyncMock()
    aeros_db_session = AsyncMock()
    project_name = "TestProject"
    simulation_duration = 8760
    overhaul_duration = 72
    overhaul_interval = 8000
    offset = 0
    schematic_name = "TestSchematic"
    
    custom_input = {
        "EQ001": {
            "mttr": 10,
            "failure_rate": 0.5,
            "distribution": "Exponential2"
        }
    }
    
    # Mock get_aeros_schematic_by_name
    mock_schematic = MagicMock()
    mock_schematic.id = 1
    mocker.patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic)
    
    # Mock db_session.execute
    mock_eq = MagicMock()
    mock_eq.location_tag = "EQ001"
    
    mock_eq_result = MagicMock()
    mock_eq_result.scalars.return_value.all.return_value = [mock_eq]
    
    mock_empty_result = MagicMock()
    mock_empty_result.scalars.return_value.all.return_value = []

    db_session.execute.side_effect = [
        mock_eq_result,    # equipments
        mock_empty_result, # rbd_group
        mock_empty_result  # non_repairable
    ]
    
    # Mock get_aeros_equipment_by_location_tag
    mock_nodes_data = [
        {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0}
    ]
    mocker.patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data)
    
    # Mock get_asset_batch (should be called but maybe ignored if custom input matches)
    mocker.patch("src.aeros_equipment.service.get_asset_batch", return_value={})
    
    mocker.patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None)
    mock_update_node = mocker.patch("src.aeros_equipment.service.update_node", return_value=None)
    
    # Run
    results = await update_equipment_for_simulation(
        db_session=db_session,
        aeros_db_session=aeros_db_session,
        project_name=project_name,
        simulation_duration=simulation_duration,
        overhaul_duration=overhaul_duration,
        overhaul_interval=overhaul_interval,
        offset=offset,
        schematic_name=schematic_name,
        custom_input=custom_input
    )
    
    assert "EQ001" in results
    assert results["EQ001"]["mttr"] == 10
    assert results["EQ001"]["distribution"] == "Exponential2"