"""InfectionWorldBuilder - extends WorldBuilder to set up infection simulation worlds."""

import sys
import random
from pathlib import Path
from typing import Dict, Any

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

from layer2_simulation.world_builder import WorldBuilder
from layer1_core.managers.population import Population
from layer1_core.managers.schedule import Schedule
from layer1_core.managers.network_manager import NetworkManager
from individual import Individual
from infection import Virus, Infection, InfectionCheckEvent
from exposure_network import create_random_network


class InfectionWorld:
    """
    Container for the infection simulation world state.

    Holds the population, schedule, virus, network, and other simulation components.
    """

    def __init__(self, population: Population, schedule: Schedule, virus: Virus, network_manager: NetworkManager):
        """
        Initialize an InfectionWorld.

        Args:
            population: Population of individuals
            schedule: Event schedule
            virus: Virus definition
            network_manager: Network of exposure relationships
        """
        self.population = population
        self.schedule = schedule
        self.virus = virus
        self.network_manager = network_manager
        self.infection_counter = 0  # Counter for generating infection IDs

    def next_infection_id(self) -> str:
        """Generate a unique infection ID."""
        self.infection_counter += 1
        return f"infection_{self.infection_counter}"


class InfectionWorldBuilder(WorldBuilder):
    """
    WorldBuilder extension for creating infection simulation worlds.

    Creates a population, virus, schedule, and initial infections.
    """

    def __init__(self, param_dict: Dict[str, Any]):
        """
        Initialize an InfectionWorldBuilder.

        Args:
            param_dict: Dictionary with keys:
                - population_size: Number of individuals
                - min_age: Minimum age of individuals
                - max_age: Maximum age of individuals
                - virus_name: Name of the virus
                - transmission_rate: Virus transmission rate (0-1)
                - incubation_period: Time steps before infectious
                - infectious_period: Time steps of infectiousness
                - recovery_period: Time steps to recover
                - base_mortality_rate: Base death probability (0-1)
                - initial_infections: Number of initial infections
                - avg_degree: Average number of contacts per individual (default: 10)
                - network_type: Type of network ("erdos_renyi", "barabasi_albert", "watts_strogatz", default: "erdos_renyi")
                - random_seed: Random seed for reproducibility (optional, inherited from WorldBuilder)
        """
        super().__init__(param_dict)

    def build(self) -> InfectionWorld:
        """
        Build the initial infection world state.

        Returns:
            InfectionWorld with population, schedule, and virus
        """
        # Create virus
        virus = Virus(
            name=self.param_dict.get('virus_name', 'DefaultVirus'),
            transmission_rate=self.param_dict.get('transmission_rate', 0.3),
            incubation_period=self.param_dict.get('incubation_period', 2),
            infectious_period=self.param_dict.get('infectious_period', 5),
            recovery_period=self.param_dict.get('recovery_period', 7),
            base_mortality_rate=self.param_dict.get('base_mortality_rate', 0.01)
        )

        # Create population
        population_size = self.param_dict.get('population_size', 100)
        min_age = self.param_dict.get('min_age', 20)
        max_age = self.param_dict.get('max_age', 80)

        population = Population(population_id="main_population")

        for i in range(population_size):
            age = random.randint(min_age, max_age)
            individual = Individual(agent_id=f"person_{i}", age=age)
            population.add_agent(individual)

        # Create schedule
        schedule = Schedule()

        # Create exposure network
        avg_degree = self.param_dict.get('avg_degree', 10)
        network_type = self.param_dict.get('network_type', 'erdos_renyi')
        network_manager = create_random_network(
            agents=population.get_agents(),
            avg_degree=avg_degree,
            network_type=network_type
        )

        # Create world container
        world = InfectionWorld(
            population=population,
            schedule=schedule,
            virus=virus,
            network_manager=network_manager
        )

        # Create initial infections
        initial_infections = self.param_dict.get('initial_infections', 1)
        agents = population.get_agents(fraction=initial_infections / population_size)

        for agent in agents:
            infection = Infection(
                agent_id=world.next_infection_id(),
                virus=virus,
                host=agent,
                time_infected=0
            )
            agent.add_infection(infection)

            # Schedule first infection check when infection becomes infectious
            first_check_time = infection.infectious_time
            check_event = InfectionCheckEvent(
                time=first_check_time,
                infection=infection,
                network_manager=network_manager,
                infection_id_generator=world.next_infection_id
            )
            schedule.add_event(step=first_check_time, event=check_event)

        self.world = world
        return world

    def get_world(self) -> InfectionWorld:
        """
        Get the constructed infection world.

        Returns:
            InfectionWorld instance
        """
        return self.world