"""Entity state collector for detailed entity tracking."""

from typing import Dict, Any, List
from .base_collector import BaseCollector
from world.objects import DefaultCell


class EntityCollector(BaseCollector):
    """Collects detailed entity state information."""

    def __init__(self, collection_interval: int = 1000, include_cells: bool = True, include_food: bool = False):
        super().__init__(collection_interval)
        self.include_cells = include_cells
        self.include_food = include_food

    def collect(self, simulation_core) -> Dict[str, Any]:
        """Collect entity states from simulation core."""
        world_state = simulation_core.get_world_state()
        entities = []

        for entity_data in simulation_core.get_entity_states():
            entity_type = entity_data['type']

            # Filter by entity type based on configuration
            if entity_type == 'cell' and not self.include_cells:
                continue
            elif entity_type == 'food' and not self.include_food:
                continue

            entities.append(entity_data)

        # Calculate additional statistics for cells
        cell_stats = {}
        if self.include_cells:
            cells = [e for e in entities if e['type'] == 'cell']
            if cells:
                energies = [c['energy'] for c in cells]
                ages = [c['age'] for c in cells]
                generations = [c['generation'] for c in cells]

                cell_stats = {
                    'avg_energy': sum(energies) / len(energies),
                    'max_energy': max(energies),
                    'min_energy': min(energies),
                    'avg_age': sum(ages) / len(ages),
                    'max_age': max(ages),
                    'avg_generation': sum(generations) / len(generations),
                    'max_generation': max(generations)
                }

        # Calculate food statistics
        food_stats = {}
        if self.include_food:
            foods = [e for e in entities if e['type'] == 'food']
            if foods:
                decays = [f['decay'] for f in foods]
                food_stats = {
                    'avg_decay': sum(decays) / len(decays),
                    'max_decay': max(decays),
                    'min_decay': min(decays),
                    'fresh_food': len([f for f in foods if f['decay'] < f['max_decay'] * 0.5])
                }

        return {
            'timestamp': simulation_core.timing.last_tick_time,
            'tick_count': world_state['tick_count'],
            'entity_count': len(entities),
            'entities': entities,
            'cell_statistics': cell_stats,
            'food_statistics': food_stats,
            'collection_type': 'entities'
        }