"""Pytest tests for Virus and Infection classes."""

import sys
from pathlib import Path

# Add paths for imports
test_dir = Path(__file__).parent
framework_src_path = test_dir.parent.parent.parent / "src"
demo_src_path = test_dir.parent / "src"
sys.path.insert(0, str(framework_src_path))
sys.path.insert(0, str(demo_src_path))

from individual import Individual
from infection import Virus, Infection, InfectionEvent, RecoveryEvent, DeathEvent, InfectionCheckEvent
from layer1_core.managers.schedule import Schedule


def test_create_virus():
    """Test creating a virus with epidemiological parameters."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    assert virus.name == "TestVirus"
    assert virus.transmission_rate == 0.3
    assert virus.incubation_period == 2
    assert virus.infectious_period == 5
    assert virus.recovery_period == 7
    assert virus.base_mortality_rate == 0.01


def test_create_infection():
    """Test creating an infection in a host individual."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)

    infection = Infection(
        agent_id="infection_1",
        virus=virus,
        host=host,
        time_infected=0
    )

    assert infection.virus == virus
    assert infection.host == host
    assert infection.time_infected == 0
    assert infection.state == 'incubating'
    assert infection.infectious_time == 2  # 0 + incubation_period
    assert infection.recovery_time == 7    # 2 + infectious_period
    assert infection.resolved_time == 14   # 7 + recovery_period


def test_infection_lifecycle():
    """Test infection state transitions over time."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)
    infection = Infection(
        agent_id="infection_1",
        virus=virus,
        host=host,
        time_infected=0
    )

    # Time 0-1: Incubating
    assert infection.state == 'incubating'
    assert infection.is_infectious(0) is False
    assert infection.is_infectious(1) is False

    # Time 2-6: Infectious
    infection.update_state(2)
    assert infection.state == 'infectious'
    assert infection.is_infectious(2) is True
    assert infection.is_infectious(6) is True

    # Time 7+: Resolved (ready for recovery or death)
    infection.update_state(7)
    assert infection.state == 'resolved'
    assert infection.is_infectious(7) is False


def test_add_infection_to_individual():
    """Test adding infections to an individual."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)
    infection = Infection(
        agent_id="infection_1",
        virus=virus,
        host=host,
        time_infected=0
    )

    # Initially no infections
    assert host.has_active_infection() is False
    assert host.get_health_state() == 'Susceptible'

    # Add infection
    host.add_infection(infection)
    assert host.has_active_infection() is True
    assert host.get_health_state() == 'Infected'
    assert len(host.infections) == 1


def test_infection_event():
    """Test InfectionEvent creates new infection."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    source_host = Individual(agent_id="person_1", age=30)
    source_infection = Infection(
        agent_id="infection_1",
        virus=virus,
        host=source_host,
        time_infected=0
    )

    target_host = Individual(agent_id="person_2", age=25)

    # Create infection event
    event = InfectionEvent(
        time=5,
        source_infection=source_infection,
        target_individual=target_host
    )

    # Execute event
    infection_counter = 1
    new_infection = event.execute(
        infection_id_generator=lambda: f"infection_{infection_counter + 1}"
    )

    assert new_infection.host == target_host
    assert new_infection.virus == virus
    assert new_infection.time_infected == 5
    assert target_host.has_active_infection() is True


def test_recovery_event():
    """Test RecoveryEvent resolves infection."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)
    infection = Infection(
        agent_id="infection_1",
        virus=virus,
        host=host,
        time_infected=0
    )
    host.add_infection(infection)

    assert host.has_active_infection() is True

    # Create and execute recovery event
    event = RecoveryEvent(time=14, infection=infection)
    event.execute()

    assert infection.state == 'recovered'
    assert host.has_active_infection() is False
    assert host.alive is True
    assert host.get_health_state() == 'Susceptible'


def test_death_event():
    """Test DeathEvent kills individual."""
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)
    infection = Infection(
        agent_id="infection_1",
        virus=virus,
        host=host,
        time_infected=0
    )
    host.add_infection(infection)

    assert host.alive is True

    # Create and execute death event
    event = DeathEvent(time=10, infection=infection, individual=host)
    event.execute()

    assert infection.state == 'fatal'
    assert host.alive is False
    assert host.get_health_state() == 'Dead'
    assert len(host.infections) == 0  # Infections cleared on death


def test_multiple_infections():
    """Test individual can have multiple infections (co-infection)."""
    virus1 = Virus(
        name="Virus1",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    virus2 = Virus(
        name="Virus2",
        transmission_rate=0.5,
        incubation_period=1,
        infectious_period=3,
        recovery_period=5,
        base_mortality_rate=0.02
    )

    host = Individual(agent_id="person_1", age=30)

    infection1 = Infection(agent_id="infection_1", virus=virus1, host=host, time_infected=0)
    infection2 = Infection(agent_id="infection_2", virus=virus2, host=host, time_infected=5)

    host.add_infection(infection1)
    host.add_infection(infection2)

    assert len(host.infections) == 2
    assert host.has_active_infection() is True
    assert host.get_health_state() == 'Infected'


def test_schedule_infection_events():
    """Test scheduling and executing infection events."""
    schedule = Schedule()
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    # Create initial infected person
    person1 = Individual(agent_id="person_1", age=30)
    infection1 = Infection(agent_id="infection_1", virus=virus, host=person1, time_infected=0)
    person1.add_infection(infection1)

    # Create susceptible person who will be infected at time 5
    person2 = Individual(agent_id="person_2", age=25)

    # Schedule infection event
    infection_event = InfectionEvent(
        time=5,
        source_infection=infection1,
        target_individual=person2
    )
    schedule.add_event(step=5, event=infection_event)

    # Initially person2 is susceptible
    assert person2.has_active_infection() is False
    assert person2.get_health_state() == 'Susceptible'

    # Execute the scheduled event
    infection_counter = 1
    event = schedule.execute_next_event()
    assert event is not None
    assert schedule.get_current_step() == 5

    new_infection = event.execute(
        infection_id_generator=lambda: f"infection_{infection_counter + 1}"
    )

    # Now person2 is infected
    assert person2.has_active_infection() is True
    assert person2.get_health_state() == 'Infected'
    assert new_infection.time_infected == 5


def test_schedule_recovery_event():
    """Test scheduling and executing recovery events."""
    schedule = Schedule()
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)
    infection = Infection(agent_id="infection_1", virus=virus, host=host, time_infected=0)
    host.add_infection(infection)

    # Schedule recovery at time 14 (after full infection cycle)
    recovery_event = RecoveryEvent(time=14, infection=infection)
    schedule.add_event(step=14, event=recovery_event)

    assert host.has_active_infection() is True

    # Execute recovery
    event = schedule.execute_next_event()
    assert event is not None
    assert schedule.get_current_step() == 14

    event.execute()

    assert infection.state == 'recovered'
    assert host.has_active_infection() is False
    assert host.alive is True


def test_schedule_death_event():
    """Test scheduling and executing death events."""
    schedule = Schedule()
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    host = Individual(agent_id="person_1", age=30)
    infection = Infection(agent_id="infection_1", virus=virus, host=host, time_infected=0)
    host.add_infection(infection)

    # Schedule death at time 10
    death_event = DeathEvent(time=10, infection=infection, individual=host)
    schedule.add_event(step=10, event=death_event)

    assert host.alive is True

    # Execute death
    event = schedule.execute_next_event()
    assert event is not None
    assert schedule.get_current_step() == 10

    event.execute()

    assert infection.state == 'fatal'
    assert host.alive is False
    assert host.get_health_state() == 'Dead'


def test_schedule_multiple_events():
    """Test scheduling and executing multiple events in order."""
    schedule = Schedule()
    virus = Virus(
        name="TestVirus",
        transmission_rate=0.3,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    # Create three people
    person1 = Individual(agent_id="person_1", age=30)
    person2 = Individual(agent_id="person_2", age=25)
    person3 = Individual(agent_id="person_3", age=40)

    # Initial infection for person1
    infection1 = Infection(agent_id="infection_1", virus=virus, host=person1, time_infected=0)
    person1.add_infection(infection1)

    # Schedule events
    # 1. Person1 infects Person2 at time 3
    event1 = InfectionEvent(time=3, source_infection=infection1, target_individual=person2)
    schedule.add_event(step=3, event=event1)

    # 2. Person2 recovers at time 20
    # (we'll create this after infection happens)

    # 3. Person3 gets infected at time 7
    event3 = InfectionEvent(time=7, source_infection=infection1, target_individual=person3)
    schedule.add_event(step=7, event=event3)

    # Execute events in order
    infection_counter = 1

    # Event at time 3: Person2 gets infected
    event = schedule.execute_next_event()
    assert schedule.get_current_step() == 3
    infection2 = event.execute(infection_id_generator=lambda: f"infection_2")
    assert person2.has_active_infection() is True

    # Event at time 7: Person3 gets infected
    event = schedule.execute_next_event()
    assert schedule.get_current_step() == 7
    infection3 = event.execute(infection_id_generator=lambda: f"infection_3")
    assert person3.has_active_infection() is True

    # Check states
    assert person1.get_health_state() == 'Infected'
    assert person2.get_health_state() == 'Infected'
    assert person3.get_health_state() == 'Infected'

def test_infection_check_event():
    """Test that InfectionCheckEvent attempts transmission."""
    from layer1_core.managers.network_manager import NetworkManager
    from layer1_core.entities.edge_agent import EdgeAgent

    virus = Virus(
        name="TestVirus",
        transmission_rate=1.0,  # 100% transmission for testing
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    # Create two individuals with a network connection
    person1 = Individual(agent_id="person_1", age=30)
    person2 = Individual(agent_id="person_2", age=25)

    network = NetworkManager()
    edge = EdgeAgent(person1, person2, "edge_1")
    network.add_edge(edge)

    # Infect person1
    infection1 = Infection(
        agent_id="infection_1",
        virus=virus,
        host=person1,
        time_infected=0
    )
    person1.add_infection(infection1)

    # Create schedule
    schedule = Schedule()

    # Create infection check event at infectious time
    check_event = InfectionCheckEvent(
        time=infection1.infectious_time,
        infection=infection1,
        network_manager=network,
        infection_id_generator=lambda: "infection_2"
    )

    # Execute the check
    new_infections = check_event.execute(schedule=schedule)

    # With 100% transmission, person2 should be infected
    assert len(new_infections) == 1
    assert person2.has_active_infection() is True


def test_transmission_scheduling():
    """Test that transmission checks schedule follow-up checks."""
    from layer1_core.managers.network_manager import NetworkManager
    from layer1_core.entities.edge_agent import EdgeAgent

    virus = Virus(
        name="TestVirus",
        transmission_rate=0.5,
        incubation_period=2,
        infectious_period=5,
        recovery_period=7,
        base_mortality_rate=0.01
    )

    person1 = Individual(agent_id="person_1", age=30)
    person2 = Individual(agent_id="person_2", age=25)

    network = NetworkManager()
    edge = EdgeAgent(person1, person2, "edge_1")
    network.add_edge(edge)

    infection1 = Infection(
        agent_id="infection_1",
        virus=virus,
        host=person1,
        time_infected=0
    )
    person1.add_infection(infection1)

    schedule = Schedule()

    # Execute check at infectious time
    current_time = infection1.infectious_time
    infection1.update_state(current_time)

    infection1.attempt_transmission(
        current_time=current_time,
        network_manager=network,
        schedule=schedule,
        infection_id_generator=lambda: "infection_2"
    )

    # Should have scheduled next check
    assert len(schedule.get_events()) > 0
