"""Infection agent and virus model classes."""

import sys
import random
from pathlib import Path
from typing import Optional, List, TYPE_CHECKING

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

from layer1_core.entities.agent import Agent

if TYPE_CHECKING:
    from layer1_core.managers.network_manager import NetworkManager
    from layer1_core.managers.schedule import Schedule


class Virus:
    """
    Represents a virus with its epidemiological parameters.

    This is not an Agent - it's a shared configuration for all infections
    of the same virus type.
    """

    def __init__(self,
                 name: str,
                 transmission_rate: float,
                 incubation_period: int,
                 infectious_period: int,
                 recovery_period: int,
                 base_mortality_rate: float):
        """
        Initialize a Virus.

        Args:
            name: Name of the virus (e.g., "SARS-CoV-2", "Influenza A")
            transmission_rate: Probability of transmission per contact with infectious individual (0-1)
            incubation_period: Time steps before infection becomes infectious
            infectious_period: Time steps during which infection is contagious
            recovery_period: Time steps until recovery (after infectious period)
            base_mortality_rate: Base probability of death (0-1), modified by age and other factors
        """
        self.name = name
        self.transmission_rate = transmission_rate
        self.incubation_period = incubation_period
        self.infectious_period = infectious_period
        self.recovery_period = recovery_period
        self.base_mortality_rate = base_mortality_rate

    def __repr__(self):
        return f"Virus({self.name}, transmission_rate={self.transmission_rate})"


class Infection(Agent):
    """
    An infection agent representing a viral infection instance within a host individual.

    Each infection has its own lifecycle: incubating -> infectious -> resolved (recovered/fatal).
    """

    def __init__(self, agent_id: str, virus: Virus, host: 'Individual', time_infected: int):
        """
        Initialize an Infection.

        Args:
            agent_id: Unique identifier for this infection instance
            virus: Virus object with epidemiological parameters
            host: Individual agent who is infected
            time_infected: Simulation time step when infection occurred
        """
        super().__init__(agent_id)
        self.virus: Virus = virus
        self.host: 'Individual' = host
        self.time_infected: int = time_infected
        self.state: str = 'incubating'  # 'incubating', 'infectious', 'recovered', 'fatal'

        # Timeline (calculated from virus parameters)
        self.infectious_time: int = time_infected + virus.incubation_period
        self.recovery_time: int = self.infectious_time + virus.infectious_period
        self.resolved_time: int = self.recovery_time + virus.recovery_period

    def is_infectious(self, current_time: int) -> bool:
        """Check if infection is currently in infectious state."""
        return (current_time >= self.infectious_time and
                current_time < self.recovery_time and
                self.state == 'infectious')

    def update_state(self, current_time: int) -> None:
        """
        Update infection state based on current time.

        Args:
            current_time: Current simulation time step
        """
        if self.state in ['recovered', 'fatal']:
            return  # Terminal states

        if current_time >= self.recovery_time:
            self.state = 'resolved'  # Will become 'recovered' or 'fatal' via events
        elif current_time >= self.infectious_time:
            self.state = 'infectious'
        else:
            self.state = 'incubating'

    def resolve_as_recovered(self) -> None:
        """Mark this infection as recovered."""
        self.state = 'recovered'

    def resolve_as_fatal(self) -> None:
        """Mark this infection as fatal."""
        self.state = 'fatal'

    def attempt_transmission(
        self,
        current_time: int,
        network_manager: 'NetworkManager',
        schedule: 'Schedule',
        infection_id : int
    ) -> List['Infection']:
        """
        Attempt to transmit infection to susceptible neighbors.

        Args:
            current_time: Current simulation time step
            network_manager: Network containing exposure relationships
            schedule: Schedule to add infection events to
            infection_id_generator: Function to generate unique infection IDs

        Returns:
            List of new Infection agents created
        """
        new_infections = []

        # Only attempt transmission if infectious
        if not self.is_infectious(current_time):
            return new_infections

        # Get neighbors from network
        neighbors = network_manager.get_neighbors(self.host)

        # Attempt to infect each susceptible neighbor
        for neighbor in neighbors:
            # Skip if already infected or dead
            if neighbor.has_active_infection() or not neighbor.alive:
                continue

            # Roll the dice: does transmission occur?
            if random.random() < self.virus.transmission_rate:
                new_infection = self.transmit(neighbor, network_manager, schedule, infection_id)
                new_infections.append(new_infection)


        # Schedule next infection check for this infection (if still infectious next step)
        next_time = current_time + 1
        if next_time < self.recovery_time:
            check_event = InfectionCheckEvent(
                time=next_time,
                infection=self,
                network_manager=network_manager,
                infection_id=infection_id
            )
            schedule.add_event(step=next_time, event=check_event)

        return new_infections

def transmit(
        self,
        target: 'Agent',
        network_manager: 'NetworkManager',
        schedule : 'Schedule',
        infection_id : int,
        **kwargs
) -> Infection :


    # Create new infection
    virus = self.source_infection.virus if self.source_infection else kwargs.get('virus')
    new_infection = Infection(
        agent_id=infection_id,
        virus=virus,
        host=self.target_individual,
        time_infected=self.time
    )
    self.target_individual.add_infection(new_infection)

    # Schedule first infection check for when new infection becomes infectious
    first_check_time = new_infection.infectious_time
    check_event = InfectionCheckEvent(
        time=first_check_time,
        infection=new_infection,
        network_manager=network_manager,
        infection_id=infection_id
    )
    schedule.add_event(step=first_check_time, event=check_event)

    return new_infection

def __repr__(self):
        return f"Infection({self.agent_id}, virus={self.virus.name}, host={self.host.agent_id}, state={self.state})"


class InfectionEvent:
    """Event representing a new infection occurring sometime in the future."""

    def __init__(self, target_individual: 'Individual', network_manager: 'NetworkManager', schedule: 'Schedule', infection_id : int) -> None:
        """
        Initialize an InfectionEvent.

        Args:
            time: Time step when infection occurs
            source_infection: Infection agent that caused this transmission (or None for initial infections)
            target_individual: Individual agent being infected
        """
        self.target_individual: 'Individual' = target_individual
        self.network_manager: 'NetworkManager' = network_manager
        self.schedule: 'Schedule' = schedule
        self.infection_id = infection_id

    def execute(self, **kwargs) -> Infection:
        """
        Execute the infection event.

        Args:
            infection_id_generator: Function to generate unique infection IDs
            **kwargs: Additional context (e.g., schedule for scheduling follow-up events)

        Returns:
            Infection: The newly created infection agent
        """
        new_infection = Infection()
        new_infection.transmit(self.target_individual, self.network_manager, self.schedule, self.infection_id, **kwargs)

   # def __repr__(self):
        #source_id = self.source_infection.agent_id if self.source_infection else "initial"
        #return f"InfectionEvent(time={self.time}, source={source_id}, target={self.target_individual.agent_id})"


class RecoveryEvent:
    """Event representing recovery from an infection."""

    def __init__(self, time: int, infection: Infection):
        """
        Initialize a RecoveryEvent.

        Args:
            time: Time step when recovery occurs
            infection: Infection agent being resolved
        """
        self.time: int = time
        self.infection: Infection = infection

    def execute(self, **kwargs) -> None:
        """Execute the recovery event."""
        self.infection.resolve_as_recovered()
        self.infection.host.remove_infection(self.infection)

    def __repr__(self):
        return f"RecoveryEvent(time={self.time}, infection={self.infection.agent_id})"


class DeathEvent:
    """Event representing death from an infection."""

    def __init__(self, time: int, infection: Infection, individual: 'Individual'):
        """
        Initialize a DeathEvent.

        Args:
            time: Time step when death occurs
            infection: Infection agent causing death
            individual: Individual agent who dies
        """
        self.time: int = time
        self.infection: Infection = infection
        self.individual: 'Individual' = individual

    def execute(self, **kwargs) -> None:
        """Execute the death event."""
        self.infection.resolve_as_fatal()
        self.individual.die()

    def __repr__(self):
        return f"DeathEvent(time={self.time}, individual={self.individual.agent_id})"


class InfectionCheckEvent:
    """Event representing a check for transmission from an infectious individual."""

    def __init__(
        self,
        time: int,
        infection: Infection,
        network_manager: 'NetworkManager',
        infection_id_generator
    ):
        """
        Initialize an InfectionCheckEvent.

        Args:
            time: Time step when check occurs
            infection: Infection agent attempting transmission
            network_manager: Network to find neighbors
            infection_id_generator: Function to generate infection IDs
        """
        self.time: int = time
        self.infection: Infection = infection
        self.network_manager = network_manager
        self.infection_id_generator = infection_id_generator

    def execute(self, **kwargs) -> List[Infection]:
        """
        Execute the infection check event.

        Returns:
            List of new infections created
        """
        # Get schedule from kwargs if provided, or use the one stored
        schedule = kwargs.get('schedule')

        # Update infection state
        self.infection.update_state(self.time)

        # Attempt transmission
        if schedule:
            return self.infection.attempt_transmission(
                current_time=self.time,
                network_manager=self.network_manager,
                schedule=schedule,
                infection_id_generator=self.infection_id_generator
            )
        return []

    def __repr__(self):
        return f"InfectionCheckEvent(time={self.time}, infection={self.infection.agent_id})"