"""
Energy Management - Token/budget tracking for The Grid.

"Programs need Energy to run on The Grid."
"""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional
from enum import Enum
import threading


class EnergyLevel(Enum):
    """Energy level thresholds."""
    FULL = "full"           # > 80%
    NORMAL = "normal"       # 40-80%
    LOW = "low"             # 10-40%
    CRITICAL = "critical"   # < 10%


@dataclass
class EnergyTransaction:
    """Record of energy consumption or allocation."""
    timestamp: datetime
    entity_id: str
    entity_type: str  # grid, cluster, block, thread, program
    amount: int
    action: str  # allocated, consumed, returned
    balance_after: int
    description: str = ""


class EnergyPool:
    """An energy pool for a Grid entity."""

    def __init__(self, budget: int, name: str, parent: Optional["EnergyPool"] = None):
        self.budget = budget
        self.name = name
        self.parent = parent
        self.consumed = 0
        self.allocated_to_children = 0
        self.transactions: list[EnergyTransaction] = []
        self._lock = threading.Lock()

    @property
    def available(self) -> int:
        """Get available energy (not consumed or allocated to children)."""
        return self.budget - self.consumed - self.allocated_to_children

    @property
    def remaining(self) -> int:
        """Get remaining energy (budget minus consumed)."""
        return self.budget - self.consumed

    @property
    def level(self) -> EnergyLevel:
        """Get current energy level."""
        pct = self.percentage
        if pct > 80:
            return EnergyLevel.FULL
        elif pct > 40:
            return EnergyLevel.NORMAL
        elif pct > 10:
            return EnergyLevel.LOW
        else:
            return EnergyLevel.CRITICAL

    @property
    def percentage(self) -> float:
        """Get energy as percentage of budget."""
        if self.budget == 0:
            return 0.0
        return (self.remaining / self.budget) * 100

    def consume(self, amount: int, entity_id: str, description: str = "") -> bool:
        """Consume energy from this pool."""
        with self._lock:
            if amount > self.available:
                return False
            self.consumed += amount
            self._record_transaction(
                entity_id=entity_id,
                entity_type="consumer",
                amount=amount,
                action="consumed",
                description=description,
            )
            return True

    def allocate_to_child(self, amount: int, child_name: str) -> Optional["EnergyPool"]:
        """Allocate energy to a child pool."""
        with self._lock:
            if amount > self.available:
                return None
            self.allocated_to_children += amount
            child_pool = EnergyPool(budget=amount, name=child_name, parent=self)
            self._record_transaction(
                entity_id=child_name,
                entity_type="child_pool",
                amount=amount,
                action="allocated",
                description=f"Allocated to {child_name}",
            )
            return child_pool

    def return_unused(self) -> int:
        """Return unused energy to parent pool."""
        if not self.parent:
            return 0
        with self._lock:
            unused = self.available
            if unused > 0:
                self.parent.allocated_to_children -= unused
                self.budget -= unused
                self._record_transaction(
                    entity_id=self.name,
                    entity_type="pool",
                    amount=unused,
                    action="returned",
                    description=f"Returned to {self.parent.name}",
                )
            return unused

    def _record_transaction(
        self,
        entity_id: str,
        entity_type: str,
        amount: int,
        action: str,
        description: str = "",
    ) -> None:
        """Record an energy transaction."""
        self.transactions.append(EnergyTransaction(
            timestamp=datetime.now(),
            entity_id=entity_id,
            entity_type=entity_type,
            amount=amount,
            action=action,
            balance_after=self.remaining,
            description=description,
        ))

    def get_bar(self, width: int = 20) -> str:
        """Get a visual energy bar."""
        filled = int((self.percentage / 100) * width)
        empty = width - filled
        return "█" * filled + "░" * empty

    def summary(self) -> str:
        """Get energy summary."""
        return (
            f"{self.name}: {self.get_bar()} "
            f"{self.remaining:,}/{self.budget:,} ({self.percentage:.0f}%)"
        )


class EnergyManager:
    """
    Manages energy (token budgets) across The Grid.

    Energy flows:
    Grid → Cluster → Block → Thread/Program

    Each level can allocate energy to children but cannot
    exceed its own budget.
    """

    def __init__(
        self,
        grid_budget: int = 10000,
        cluster_budget: int = 2000,
        block_budget: int = 500,
        thread_budget: int = 100,
        spawn_cost: int = 100,
        low_threshold: float = 0.1,
    ):
        self.grid_budget = grid_budget
        self.cluster_budget = cluster_budget
        self.block_budget = block_budget
        self.thread_budget = thread_budget
        self.spawn_cost = spawn_cost
        self.low_threshold = low_threshold

        # Root energy pool
        self.grid_pool = EnergyPool(budget=grid_budget, name="Grid")

        # Track all pools
        self.pools: dict[str, EnergyPool] = {"grid": self.grid_pool}

    def create_cluster_pool(self, cluster_id: str) -> Optional[EnergyPool]:
        """Create an energy pool for a new Cluster."""
        pool = self.grid_pool.allocate_to_child(
            amount=min(self.cluster_budget, self.grid_pool.available),
            child_name=f"cluster:{cluster_id}",
        )
        if pool:
            self.pools[f"cluster:{cluster_id}"] = pool
        return pool

    def create_block_pool(self, cluster_id: str, block_id: str) -> Optional[EnergyPool]:
        """Create an energy pool for a new Block."""
        cluster_pool = self.pools.get(f"cluster:{cluster_id}")
        if not cluster_pool:
            return None
        pool = cluster_pool.allocate_to_child(
            amount=min(self.block_budget, cluster_pool.available),
            child_name=f"block:{block_id}",
        )
        if pool:
            self.pools[f"block:{block_id}"] = pool
        return pool

    def create_thread_pool(self, block_id: str, thread_id: str) -> Optional[EnergyPool]:
        """Create an energy pool for a new Thread."""
        block_pool = self.pools.get(f"block:{block_id}")
        if not block_pool:
            return None
        pool = block_pool.allocate_to_child(
            amount=min(self.thread_budget, block_pool.available),
            child_name=f"thread:{thread_id}",
        )
        if pool:
            self.pools[f"thread:{thread_id}"] = pool
        return pool

    def consume(self, pool_id: str, amount: int, description: str = "") -> bool:
        """Consume energy from a pool."""
        pool = self.pools.get(pool_id)
        if not pool:
            return False
        return pool.consume(amount, pool_id, description)

    def get_pool(self, pool_id: str) -> Optional[EnergyPool]:
        """Get a pool by ID."""
        return self.pools.get(pool_id)

    def is_low(self, pool_id: str) -> bool:
        """Check if a pool is at low energy."""
        pool = self.pools.get(pool_id)
        if not pool:
            return True
        return pool.percentage <= (self.low_threshold * 100)

    def grid_summary(self) -> str:
        """Get full grid energy summary."""
        lines = ["ENERGY FLOW", "═" * 40]
        lines.append(self.grid_pool.summary())

        # Show cluster pools
        for pool_id, pool in self.pools.items():
            if pool_id.startswith("cluster:"):
                lines.append(f"  └─ {pool.summary()}")
                # Show block pools under this cluster
                for block_id, block_pool in self.pools.items():
                    if block_id.startswith("block:") and block_pool.parent == pool:
                        lines.append(f"       ├─ {block_pool.summary()}")

        return "\n".join(lines)

    def total_consumed(self) -> int:
        """Get total energy consumed across all pools."""
        return sum(pool.consumed for pool in self.pools.values())

    def total_remaining(self) -> int:
        """Get total remaining energy in grid."""
        return self.grid_pool.remaining