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():
    # 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
    
    # Mock db_session.execute results
    mock_eq = MagicMock()
    mock_eq.location_tag = "EQ001"
    mock_eq_result = MagicMock()
    mock_eq_result.scalars.return_value.all.return_value = [mock_eq]
    
    mock_group = MagicMock()
    mock_group.group_name = "GROUP001"
    mock_group_result = MagicMock()
    mock_group_result.scalars.return_value.all.return_value = [mock_group]
    
    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 nodes data
    mock_nodes_data = [
        {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0},
        {"equipmentName": "GROUP001", "relDisType": "Fixed", "relDisP1": 0}
    ]
    
    # Mock reliability data
    mock_reliability_data = {
        "EQ001": {
            "relDisType": "Weibull2",
            "relDisP1": 1.5,
            "relDisP2": 10000,
            "cmDisType": "Fixed",
            "cmDisP1": 5,
            "cmDisP2": 0
        }
    }

    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
         patch("src.aeros_equipment.service.get_asset_batch", return_value=mock_reliability_data), \
         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
         patch("src.aeros_equipment.service.update_node", return_value=None):
        
        # Run the function
        res = 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
        )
        
        results = res['reliability']
        equipment_nodes = res['equipment_nodes']
        
        # Assertions
        assert "EQ001" in results
        assert results["EQ001"]["distribution"] == "Weibull2"
        assert results["EQ001"]["oh_duration"] == overhaul_duration
        
        # 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():
    # 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
    
    # 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 nodes data
    mock_nodes_data = [
        {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0}
    ]
    
    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
         patch("src.aeros_equipment.service.get_asset_batch", return_value={}), \
         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
         patch("src.aeros_equipment.service.update_node", return_value=None):
        
        # Run
        res = 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=8000,
            offset=0,
            schematic_name=schematic_name,
            custom_input=custom_input
        )
        
        results = res['reliability']
        
        assert "EQ001" in results
        assert results["EQ001"]["mttr"] == 10
        assert results["EQ001"]["distribution"] == "Exponential2"

@pytest.mark.asyncio
async def test_update_equipment_for_simulation_missing_reliability():
    # 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
    
    # Mock db_session.execute
    mock_eq = MagicMock()
    mock_eq.location_tag = "EQ_MISSING"
    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 nodes data
    mock_nodes_data = [
        {"equipmentName": "EQ_MISSING", "relDisType": "Fixed", "relDisP1": 0}
    ]
    
    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
         patch("src.aeros_equipment.service.get_asset_batch", return_value={}), \
         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
         patch("src.aeros_equipment.service.update_node", return_value=None):
        
        # Run
        res = 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
        )
        
        results = res['reliability']
        equipment_nodes = res['equipment_nodes']
        
        assert "EQ_MISSING" in results
        # Default values when no reliability data
        assert results["EQ_MISSING"]["mttr"] == 0
        assert results["EQ_MISSING"]["distribution"] == "Fixed"
        
        # Check defaults in returned equipment_nodes
        eq_node = equipment_nodes[0]
        assert eq_node["cmDisUnitCode"] == "UHour"
        assert eq_node["relDisUnitCode"] == "UHour"
        assert eq_node["ohDisUnitCode"] == "UHour"

@pytest.mark.asyncio
async def test_update_equipment_for_simulation_sigma_constraint():
    # 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
    
    # Mock db_session.execute
    mock_eq = MagicMock()
    mock_eq.location_tag = "EQ_SIGMA"
    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 nodes data
    mock_nodes_data = [
        {"equipmentName": "EQ_SIGMA", "relDisType": "Fixed", "relDisP1": 0}
    ]
    
    # Mock reliability data with high sigma (> 1000)
    mock_reliability_data = {
        "EQ_SIGMA": {
            "relDisType": "Normal",
            "relDisP1": 100,  # mu
            "relDisP2": 1500, # sigma
            "cmDisType": "Fixed",
            "cmDisP1": 0,
            "cmDisP2": 0
        }
    }
    
    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
         patch("src.aeros_equipment.service.get_asset_batch", return_value=mock_reliability_data), \
         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
         patch("src.aeros_equipment.service.update_node", return_value=None):
        
        # Run
        res = 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
        )
        
        # Check if sigma was constrained correctly in returned equipment_nodes
        equipment_nodes = res['equipment_nodes']
        eq_node = equipment_nodes[0]
        assert eq_node["relDisP2"] == 999