from hiveos.models.chunk import Chunk
from hiveos.models.task import Task
from hiveos.models.agent import Agent
from hiveos.utilities.render_zone_summary import render_zone_summary

class ZoneController:
    def __init__(self, zone_id):
        self.zone_id = zone_id
        self.task_id = None
        self.required_capability = 'generic'
        self.agent_ids = []
        self.agent_instances = []
        self.should_shutdown = False
        self.needs_more_agents = False
        self.fallback_buffer = 3

    def assign_task(self, task_id):
        self.task_id = task_id
        
        if not self.agent_instances:
            self.needs_more_agents = True

        chunks = Chunk.get_pending_by_task(task_id)
        if chunks:
            self.required_capability = chunks[0][2]  # capability from the chunk tuple
        else:
            self.required_capability = 'generic'

    def assign_agents(self, agent_instances, append=False):
        if append:
            new_agents = [a for a in agent_instances if a.agent_id not in self.agent_ids]
            self.agent_instances += new_agents
            self.agent_ids += [a.agent_id for a in new_agents]
        else:
            self.agent_instances = agent_instances
            self.agent_ids = [a.agent_id for a in agent_instances]

        for agent in agent_instances:
            agent.zone = self.zone_id
            Agent.lock_to_zone(agent.agent_id, self.zone_id)

    def release_agent(self, agent_id):
        self.agent_instances = [a for a in self.agent_instances if a.agent_id != agent_id]
        self.agent_ids = [id for id in self.agent_ids if id != agent_id]

    def agents_needed(self):
        pending_chunks = Chunk.get_pending_by_task(self.task_id)
        idle_agents = [a for a in self.agent_instances if a.status == 'idle']
        need = max(len(pending_chunks) - len(idle_agents), 0)
        return need + self.fallback_buffer

    def releasable_agents(self):
        pending_chunks = Chunk.get_pending_by_task(self.task_id)
        idle_agents = [a for a in self.agent_instances if a.status == 'idle']
        target = max(len(pending_chunks) + self.fallback_buffer, 0)
        excess = len(idle_agents) - target
        if excess <= 0:
            return []
        return idle_agents[:excess]

    def tick(self, verbose=False):
        
        for agent in self.agent_instances:
            agent.tick()
            if verbose:
                print(f"[Zone {self.zone_id}] Agent {agent.agent_id} status: {agent.status}")
        
        chunks = Chunk.get_pending_by_task(self.task_id)
        
        """
        chunks = [
            c for c in Chunk.get_pending_by_task(self.task_id)
            if not c['depends_on'] or Chunk.get_status(c['depends_on']) == 'complete'
        ]
        """
        
        available_agents = [a for a in self.agent_instances
                            if a.status == 'idle' and a.current_chunk is None]

    
        for chunk_row in chunks:
            chunk_id = chunk_row[0]
            

            if not available_agents:
                self.needs_more_agents = True
                break
                

            agent = available_agents.pop(0)
            agent.assign_chunk(chunk_id)
            agent.status = 'working'
            Agent.update_status(agent.agent_id, 'working')

        

        if not verbose:
            self.print_summary()

        if Chunk.all_chunks_completed(self.task_id):
            self.should_shutdown = True

    def sync(self):
        """
        TODO: Zone Heartbeat Sync

        This method will eventually push the runtime state of all agents in this zone
        to the central database (or monitoring system) for diagnostic and monitoring purposes.

        Optimization idea:
        - Track a 'dirty' flag on each agent (set True when status, battery, zone, or chunk changes)
        - Only sync agents where agent.dirty == True
        - Reset the flag after sync
        - Use this for efficient state propagation to external interfaces or APIs
        - Could also push to a separate `agent_snapshots` table for historical analysis

        Triggered by: core.tick() → zone.sync()
        """
        pass
        

    """
    def print_summary(self):
        total = len(self.agent_instances)
        idle = sum(1 for a in self.agent_instances if a.status == 'idle')
        working = sum(1 for a in self.agent_instances if a.status == 'working')
        recharge = sum(1 for a in self.agent_instances if a.status == 'recharging')
        cooldown = sum(1 for a in self.agent_instances if a.status == 'cooldown')

        total_chunks = Task.chunk_count(self.task_id)
        pending_chunks = len(Chunk.get_pending_by_task(self.task_id))
        completed_chunks = Chunk.count_completed(self.task_id)

        print(f"Zone {self.zone_id} | Agents: {total} (Working: {working}, Idle: {idle}, Recharge: {recharge}, Cooldown: {cooldown})")
        print(f"           Chunks: {total_chunks} total, {completed_chunks} completed, {pending_chunks} pending")
        
    """
    
    def print_summary(self):
        
        total_chunks = Task.chunk_count(self.task_id)
        pending_chunks = len(Chunk.get_pending_by_task(self.task_id))
        completed_chunks = Chunk.count_completed(self.task_id)
        
        summary = render_zone_summary(
            self.zone_id,
            Task.get_by_id(self.task_id).name,
            total_chunks,
            completed_chunks,
            {
                "total": len(self.agent_instances),
                "working": sum(a.status == 'working' for a in self.agent_instances),
                "idle": sum(a.status == 'idle' for a in self.agent_instances),
                "recharging": sum(a.status == 'recharging' for a in self.agent_instances),
                "cooldown": sum(a.status == 'cooldown' for a in self.agent_instances),
            }
        )
        print(summary)

        """
        USE_COLORS = False  # Flip to True for ANSI-compatible terminals

        def c(text, color_code):
            return f"{color_code}{text}\033[0m" if USE_COLORS else text

        GREEN = "\033[92m"
        YELLOW = "\033[93m"
        RED = "\033[91m"

        total = len(self.agent_instances)
        working = sum(1 for a in self.agent_instances if a.status == "WORKING")
        idle = sum(1 for a in self.agent_instances if a.status == "IDLE")
        recharge = sum(1 for a in self.agent_instances if a.status == "RECHARGING")
        cooldown = sum(1 for a in self.agent_instances if a.status == "COOLDOWN")

        task_name = "(none)"
        capability = self.required_capability
        if self.task_id:
            task = Task.get_by_id(self.task_id)
            if task:
                task_name = task[1]

            total_chunks = Task.chunk_count(self.task_id)
            completed_chunks = Chunk.count_completed(self.task_id)
            pending_chunks = len(Chunk.get_pending_by_task(self.task_id))
        else:
            total_chunks = completed_chunks = pending_chunks = 0

        # Progress bar
        bar_length = 20
        progress = completed_chunks / total_chunks if total_chunks else 0
        bar = "█" * int(bar_length * progress) + "-" * (bar_length - int(bar_length * progress))

        print(f"\n{c('🌐 Zone ' + str(self.zone_id), GREEN)} | Task: '{task_name}' | "
              f"Capability: {capability} | Agents: {total} "
              f"(🟩 {working}, 🟨 {idle}, 🔋 {recharge}, ❄️ {cooldown})")

        print(f"📦 Chunks: {completed_chunks}/{total_chunks} Complete [{bar}] {int(progress * 100)}%")

        if completed_chunks == total_chunks and total_chunks > 0:
            print(c(f"✅ All chunks completed for task '{task_name}' in Zone {self.zone_id}", GREEN))
        elif idle > 0 and pending_chunks > 0:
            print(c("⚠️  Idle agents available while chunks remain!", YELLOW))
        elif total_chunks == 0:
            print(c(f"🚫 No chunks assigned yet in Zone {self.zone_id}.", RED))
            
        """