"""
STARWASTE - Server Entities
Entity management (players, NPCs, enemies, items)
"""

from typing import Dict, Optional, Tuple, TYPE_CHECKING
from dataclasses import dataclass, field
import threading
from shared.constants import ENTITY_TYPES, MAX_HEALTH, MAX_HUNGER
from shared.inventory import Inventory
from shared.utils import console

if TYPE_CHECKING:
    from server.database import Database


@dataclass
class Entity:
    """Base entity"""
    id: int
    entity_type: int
    x: float
    y: float
    vx: float = 0.0
    vy: float = 0.0

    # Persistence tracking (dirty/revision). Not part of network serialization.
    _persist_rev: int = field(default=0, init=False, repr=False)
    _persist_saved_rev: int = field(default=0, init=False, repr=False)

    def mark_dirty(self):
        self._persist_rev += 1

    def mark_clean_if_unchanged(self, rev: int) -> bool:
        if self._persist_rev == rev:
            self._persist_saved_rev = rev
            return True
        return False
    
    def update(self, dt: float):
        """Update entity (override in subclasses)"""
        pass
    
    def to_dict(self) -> dict:
        """Serialize to dictionary"""
        return {
            "id": self.id,
            "type": self.entity_type,
            "x": self.x,
            "y": self.y,
            "vx": self.vx,
            "vy": self.vy,
        }


class Player(Entity):
    """Player entity"""
    
    def __init__(self, id: int, username: str, x: float, y: float):
        super().__init__(id, ENTITY_TYPES["PLAYER"], x, y)
        self.username = username
        self.health = MAX_HEALTH
        self.hunger = MAX_HUNGER
        self.radiation = 0
        self.cold = 0  # Cold exposure (frozen biome)
        self.inventory = Inventory(main_slots=40)
        
        # Harvesting state
        self.harvesting_target = None  # Entity ID of resource being harvested
        self.harvest_progress = 0.0    # Current harvest progress (0-1)
        
    def update(self, dt: float):
        """Update player state"""
        # Update position
        self.x += self.vx * dt
        self.y += self.vy * dt
        
        # Decay hunger
        self.hunger = max(0, self.hunger - 0.1 * dt)
        
        # Take damage from hunger
        if self.hunger == 0:
            self.health = max(0, self.health - 1 * dt)
    
    def to_dict(self) -> dict:
        """Serialize player data"""
        data = super().to_dict()
        data.update({
            "username": self.username,
            "health": self.health,
            "hunger": self.hunger,
            "radiation": int(self.radiation),  # Send as integer
            "cold": int(self.cold),  # Send as integer
            "inventory": self.inventory.to_dict(),
            "harvesting_target": self.harvesting_target,
            "harvest_progress": self.harvest_progress,
        })
        return data


class Enemy(Entity):
    """Enemy entity"""
    
    def __init__(self, id: int, enemy_type: str, x: float, y: float):
        super().__init__(id, ENTITY_TYPES["ENEMY"], x, y)
        self.enemy_type = enemy_type
        self.health = 50
        self.damage = 10
        self.speed = 100
        
    def update(self, dt: float):
        """Update enemy AI"""
        # Simple movement for now
        self.x += self.vx * dt
        self.y += self.vy * dt
    
    def to_dict(self) -> dict:
        data = super().to_dict()
        data.update({
            "enemy_type": self.enemy_type,
            "health": self.health,
        })
        return data


class Item(Entity):
    """Item entity (dropped in world)"""

    def __init__(self, id: int, item_id: str = None, x: float = 0.0, y: float = 0.0, *, item_type: str = None):
        super().__init__(id, ENTITY_TYPES["ITEM"], x, y)
        # Persisted key is item_id. item_type is kept for backward compatibility.
        self.item_id = item_id or item_type or ""
        self.quantity = 1

    def set_quantity(self, quantity: int):
        new_qty = max(1, int(quantity or 1))
        if new_qty != self.quantity:
            self.quantity = new_qty
            self.mark_dirty()

    @property
    def item_type(self) -> str:
        """Backward-compatible alias."""
        return self.item_id

    @item_type.setter
    def item_type(self, value: str):
        self.item_id = value

    def to_dict(self) -> dict:
        data = super().to_dict()
        data.update({
            "item_id": self.item_id,
            # Keep old field name so older clients/tools don't break.
            "item_type": self.item_id,
            "quantity": self.quantity,
        })
        return data


class ResourceNode(Entity):
    """Resource node entity (tree, rock, scrap pile, etc.)"""
    
    # Resource node types with their properties
    RESOURCE_TYPES = {
        "tree": {
            "health": 100,
            "drops": [("resource_wood", 5, 10)],  # (item_id, min_qty, max_qty)
            "required_tool": "tool_axe",
            "harvest_time": 3.0,  # seconds
        },
        "rock": {
            "health": 150,
            "drops": [("resource_stone", 3, 8), ("resource_scrap", 1, 3)],
            "required_tool": "tool_pickaxe",
            "harvest_time": 4.0,
        },
        "scrap_pile": {
            "health": 80,
            "drops": [("resource_scrap", 10, 20), ("resource_circuits", 1, 3)],
            "required_tool": None,  # No tool required
            "harvest_time": 2.5,
        },
        "iron_ore": {
            "health": 200,
            "drops": [("resource_metal", 15, 25)],
            "required_tool": "tool_pickaxe",
            "harvest_time": 5.0,
        },
    }
    
    def __init__(self, id: int, resource_type: str, x: float, y: float, variant_id: Optional[int] = None):
        super().__init__(id, ENTITY_TYPES.get("RESOURCE", 4), x, y)
        self.resource_type = resource_type
        self.variant_id = variant_id  # Unified ID for all resources (tree_01 = 1, rock_05 = 5, etc.)
        
        # Get properties for this resource type
        props = self.RESOURCE_TYPES.get(resource_type, self.RESOURCE_TYPES["scrap_pile"])
        self.max_health = props["health"]
        self.health = self.max_health
        self.drops = props["drops"]
        self.required_tool = props["required_tool"]
        self.harvest_time = props["harvest_time"]
        
        # Harvesting state
        self.being_harvested = False
        self.harvest_progress = 0.0
        self.harvester_id: Optional[int] = None
        
    def take_damage(self, damage: float) -> bool:
        """
        Deal damage to resource node.
        Returns True if node is destroyed.
        """
        self.health -= damage
        self.mark_dirty()
        if self.health <= 0:
            self.health = 0
            return True
        return False
    
    def start_harvest(self, *, harvester_id: Optional[int] = None):
        """Start harvesting this node."""
        self.being_harvested = True
        self.harvest_progress = 0.0
        self.harvester_id = int(harvester_id) if harvester_id is not None else None
    
    def update_harvest(self, dt: float) -> bool:
        """
        Update harvest progress.
        Returns True if harvest is complete.
        """
        if not self.being_harvested:
            return False
        
        self.harvest_progress += dt
        if self.harvest_progress >= self.harvest_time:
            self.being_harvested = False
            return True
        return False
    
    def cancel_harvest(self):
        """Cancel ongoing harvest"""
        self.being_harvested = False
        self.harvest_progress = 0.0
        self.harvester_id = None
    
    def get_drops(self) -> list[Tuple[str, int]]:
        """
        Get random drops when node is destroyed.
        Returns list of (item_id, quantity) tuples.
        """
        import random
        drops = []
        for item_id, min_qty, max_qty in self.drops:
            qty = random.randint(min_qty, max_qty)
            drops.append((item_id, qty))
        return drops
    
    def to_dict(self) -> dict:
        data = super().to_dict()
        data.update({
            "resource_type": self.resource_type,
            "health": self.health,
            "max_health": self.max_health,
            "being_harvested": self.being_harvested,
            "harvest_progress": self.harvest_progress / self.harvest_time if self.harvest_time > 0 else 0,
            "harvester_id": self.harvester_id,
        })
        # Add variant ID (unified for all resource types)
        if self.variant_id is not None:
            data["variant_id"] = self.variant_id
        return data


class EntityManager:
    """Manages all entities in the world"""
    
    def __init__(self, database: Optional['Database'] = None):
        self.entities: Dict[int, Entity] = {}
        self.next_id = 1
        self.database = database

        # Dirty tracking for persistence
        self._dirty_entity_ids: set[int] = set()
        self._dirty_entities_lock = threading.Lock()
        
        # 🔥 CACHED PLAYERS: Avoid expensive isinstance checks in hot paths
        self._players: Dict[int, Player] = {}  # player_id -> Player
        
        # 🔥 SPATIAL HASH GRID for fast entity lookups
        # Grid cell size = 512px (16 tiles × 32px) - good balance
        self.grid_size = 512
        self.spatial_grid: Dict[Tuple[int, int], set[int]] = {}  # (grid_x, grid_y) -> {entity_ids}
        
    def add_entity(self, entity: Entity) -> int:
        """Add entity to world"""
        entity.id = self.next_id
        self.entities[entity.id] = entity
        self.next_id += 1
        
        # 🔥 Add to spatial grid
        self._add_to_grid(entity.id, entity.x, entity.y)
        
        # 🔥 CACHE: Update players cache
        if isinstance(entity, Player):
            self._players[entity.id] = entity
        else:
            # New non-player entities must be persisted.
            entity.mark_dirty()
            with self._dirty_entities_lock:
                self._dirty_entity_ids.add(entity.id)
        
        return entity.id

    def mark_entity_dirty(self, entity_id: int):
        entity = self.entities.get(entity_id)
        if not entity or isinstance(entity, Player):
            return
        entity.mark_dirty()
        with self._dirty_entities_lock:
            self._dirty_entity_ids.add(entity_id)
    
    def remove_entity(self, entity_id: int):
        """Remove entity from world and database"""
        if entity_id in self.entities:
            entity = self.entities[entity_id]
            
            # 🔥 Remove from spatial grid
            self._remove_from_grid(entity_id, entity.x, entity.y)
            
            # 🔥 CACHE: Update players cache
            if entity_id in self._players:
                del self._players[entity_id]
            
            # Remove from memory
            del self.entities[entity_id]

            # Keep dirty set in sync
            with self._dirty_entities_lock:
                self._dirty_entity_ids.discard(entity_id)
            
            # Remove from database (if database is connected)
            if self.database:
                self.database.delete_entity(entity_id)
                # Only log for debugging purposes (can be disabled later)
                # console.print(f"🗑️ Removed entity {entity_id} from memory and database")
    
    def get_entity(self, entity_id: int) -> Optional[Entity]:
        """Get entity by ID"""
        return self.entities.get(entity_id)
    
    def update(self, dt: float):
        """Update all entities and their spatial grid positions"""
        for entity in list(self.entities.values()):
            # Static entities have no per-tick simulation today.
            # Keeping them out of the hot update path prevents O(N) scaling issues
            # when the world accumulates many items/resources.
            if isinstance(entity, (Item, ResourceNode)):
                continue

            # Store old position for grid update check
            old_x, old_y = entity.x, entity.y
            old_cell = self._get_grid_cell(old_x, old_y)
            
            # Update entity
            entity.update(dt)
            
            # 🔥 OPTIMIZED: Only update grid if entity moved to different cell
            # This avoids expensive grid updates for small movements
            new_cell = self._get_grid_cell(entity.x, entity.y)
            if old_cell != new_cell:
                self._remove_from_grid(entity.id, old_x, old_y)
                self._add_to_grid(entity.id, entity.x, entity.y)
    
    def get_all_entities_data(self) -> list[dict]:
        """Get serialized data for all entities"""
        return [entity.to_dict() for entity in self.entities.values()]
    
    def _get_grid_cell(self, x: float, y: float) -> Tuple[int, int]:
        """Get grid cell coordinates for world position"""
        return (int(x // self.grid_size), int(y // self.grid_size))
    
    def _add_to_grid(self, entity_id: int, x: float, y: float):
        """Add entity to spatial grid"""
        cell = self._get_grid_cell(x, y)
        if cell not in self.spatial_grid:
            self.spatial_grid[cell] = set()
        self.spatial_grid[cell].add(entity_id)
    
    def _remove_from_grid(self, entity_id: int, x: float, y: float):
        """Remove entity from spatial grid"""
        cell = self._get_grid_cell(x, y)
        if cell in self.spatial_grid:
            self.spatial_grid[cell].discard(entity_id)
            # Clean up empty cells
            if not self.spatial_grid[cell]:
                del self.spatial_grid[cell]
    
    def get_entities_near(self, x: float, y: float, radius: float, entity_type: Optional[int] = None) -> list[Entity]:
        """
        🔥 OPTIMIZED: Get entities within radius using spatial grid.
        
        Args:
            x: Center X position
            y: Center Y position
            radius: Search radius in pixels
            entity_type: Optional filter by entity type (e.g., ENTITY_TYPES["TREE"])
        
        Returns:
            List of nearby entities
        """
        # Calculate which grid cells to check
        cells_to_check = int(radius // self.grid_size) + 1
        center_cell = self._get_grid_cell(x, y)
        
        nearby_entities = []
        radius_sq = radius * radius
        
        # Check surrounding grid cells
        for dx in range(-cells_to_check, cells_to_check + 1):
            for dy in range(-cells_to_check, cells_to_check + 1):
                cell = (center_cell[0] + dx, center_cell[1] + dy)
                
                if cell not in self.spatial_grid:
                    continue
                
                # Check all entities in this cell
                for entity_id in self.spatial_grid[cell]:
                    if entity_id not in self.entities:
                        continue
                    
                    entity = self.entities[entity_id]
                    
                    # Filter by type if specified
                    if entity_type is not None and entity.entity_type != entity_type:
                        continue
                    
                    # Distance check (squared)
                    dx_entity = entity.x - x
                    dy_entity = entity.y - y
                    if dx_entity * dx_entity + dy_entity * dy_entity <= radius_sq:
                        nearby_entities.append(entity)
        
        return nearby_entities
    
    def load_entities_from_db(self):
        """Load all entities from database"""
        if not self.database:
            return
        
        entities_data = self.database.load_entities()
        loaded_count = 0
        
        for data in entities_data:
            entity_type = data["entity_type"]
            entity_id = data["entity_id"]
            x = data["x"]
            y = data["y"]
            entity_data = data["data"]
            
            entity = None
            
            # Reconstruct entity based on type
            if entity_type == "resource":
                resource_type = entity_data["resource_type"]
                
                # Load variant_id from database (or regenerate if missing)
                variant_id = entity_data.get("variant_id")
                
                # Fallback: Regenerate variant_id if not in DB (backward compatibility)
                if variant_id is None and resource_type in ["tree", "rock"] and hasattr(self, '_world_seed'):
                    from shared.constants import TILE_SIZE
                    world_x = int((x - TILE_SIZE // 2) / TILE_SIZE)
                    world_y = int((y - TILE_SIZE // 2) / TILE_SIZE)
                    variant_id = (hash((world_x, world_y, self._world_seed)) % 10) + 1
                
                entity = ResourceNode(
                    id=entity_id,
                    resource_type=resource_type,
                    x=x,
                    y=y,
                    variant_id=variant_id  # Pass unified variant_id
                )
                entity.health = entity_data.get("health", entity.max_health)
                
            elif entity_type == "item":
                entity = Item(
                    id=entity_id,
                    item_id=entity_data["item_id"],
                    x=x,
                    y=y
                )
                entity.quantity = entity_data.get("quantity", 1)
            
            # Add entity (use existing ID, don't assign new one)
            if entity:
                self.entities[entity_id] = entity
                # 🔥 FIX: Add to spatial grid when loading from DB!
                self._add_to_grid(entity_id, x, y)
                self.next_id = max(self.next_id, entity_id + 1)
                # Loaded entities are clean.
                entity._persist_rev = 0
                entity._persist_saved_rev = 0
                loaded_count += 1
        
        if loaded_count > 0:
            console.print(f"💾 Loaded {loaded_count} entities from database")
    
    def save_entity(self, entity_id: int):
        """Save single entity to database"""
        if not self.database:
            return
        
        entity = self.entities.get(entity_id)
        if not entity:
            return
        
        # Don't save players (they have separate table)
        if isinstance(entity, Player):
            return
        
        # Determine entity type
        entity_type = "unknown"
        entity_data = {}
        
        if isinstance(entity, ResourceNode):
            entity_type = "resource"
            entity_data = {
                "resource_type": entity.resource_type,
                "health": entity.health,
                "max_health": entity.max_health,
                "variant_id": entity.variant_id,  # CRITICAL: Save variant_id!
            }
        elif isinstance(entity, Item):
            entity_type = "item"
            entity_data = {
                "item_id": entity.item_id,
                "quantity": entity.quantity,
            }
        
        self.database.save_entity(entity.id, entity_type, entity.x, entity.y, entity_data)

        # Mark clean if entity unchanged since we captured the rev.
        rev = getattr(entity, "_persist_rev", 0)
        entity._persist_saved_rev = rev
        with self._dirty_entities_lock:
            self._dirty_entity_ids.discard(entity_id)
    
    def save_all_entities(self, *, force: bool = False) -> int:
        """Save dirty entities to database (or all entities with force=True)."""
        if not self.database:
            return 0

        if force:
            entity_ids_to_save = [e.id for e in list(self.entities.values()) if not isinstance(e, Player)]
        else:
            with self._dirty_entities_lock:
                entity_ids_to_save = list(self._dirty_entity_ids)

        if not entity_ids_to_save:
            return 0

        saved_count = 0
        cleaned_ids: list[int] = []

        for entity_id in entity_ids_to_save:
            entity = self.entities.get(entity_id)
            if not entity or isinstance(entity, Player):
                cleaned_ids.append(entity_id)
                continue

            rev = getattr(entity, "_persist_rev", 0)
            saved_rev = getattr(entity, "_persist_saved_rev", 0)
            if rev <= saved_rev and not force:
                cleaned_ids.append(entity_id)
                continue

            self.save_entity(entity_id)
            saved_count += 1
            if entity.mark_clean_if_unchanged(rev):
                cleaned_ids.append(entity_id)

        if cleaned_ids and not force:
            with self._dirty_entities_lock:
                for entity_id in cleaned_ids:
                    self._dirty_entity_ids.discard(entity_id)

        if saved_count > 0:
            console.print(f"💾 Saved {saved_count} dirty entities to database")

        return saved_count
