from hiveos.models.chunk import Chunk
from hiveos.models.task import Task
from hiveos.models.agent import Agent
from hiveos.runtime.agent_runtime import get_idle_agents
from hiveos.runtime.chunk_runtime import chunk_objects, get_by_id as get_chunk_by_id, count_completed
from hiveos.runtime.task_runtime import task_objects, get_by_id as get_task_by_id
from hiveos.utilities.render_zone_summary import render_zone_summary
from hiveos.utilities.gate_evaluator import GateEvaluator

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
        self.gated_agents = []
        self.gate_evaluator = GateEvaluator(self)

    def assign_task(self, task_id):
        self.task_id = task_id
        if not self.agent_instances:
            self.needs_more_agents = True
        chunks = [
            c for c in chunk_objects.values()
            if c.task_id == task_id and c.status == 'pending'
        ]
        self.required_capability = chunks[0].required_capability if chunks else '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 = [
            c for c in chunk_objects.values()
            if c.task_id == self.task_id and c.status == 'pending'
        ]
        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 = [
            c for c in chunk_objects.values()
            if c.task_id == self.task_id and c.status == 'pending'
        ]
        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):
        # Step 1: Lift any gates if dependencies are now satisfied
        for agent in self.gated_agents[:]:
            if agent.status != "working" or not agent.current_chunk:
                continue
                
            chunk = get_chunk_by_id(agent.current_chunk)
            if not chunk:
                #print(f"[GATE CHECK FAIL] Could not load chunk for {agent.agent_id}")
                continue

            dep_id = chunk.depends_on
            if not dep_id:
                #print(f"[GATE CHECK] {chunk['chunk_id']} has no dependency")
                agent.ready_to_execute = True
                self.gated_agents.remove(agent)
                continue

            dep_chunk = get_chunk_by_id(dep_id)
            if not dep_chunk or dep_chunk.status != "completed":
                #print(f"[GATE CHECK FAIL] Dependency chunk {dep_id} not found")
                continue

            #print(f"[GATE LIFTED] {agent.agent_id} is cleared to execute {chunk['chunk_id']}")
            agent.ready_to_execute = True
            self.gated_agents.remove(agent)

        # Step 2: Tick all agents after lifting gates
        for agent in self.agent_instances:
            agent.tick()
            if verbose:
                print(f"[Zone {self.zone_id}] Agent {agent.agent_id} status: {agent.status}")

        # Step 3: Assign available chunks to available idle agents
        chunks = [
            c for c in chunk_objects.values()
            if c.task_id == self.task_id and c.status == 'pending'
        ]
        available_agents = [a for a in self.agent_instances if a.status == 'idle' and a.current_chunk is None]
        
        for chunk in chunks:
            compatible_agents = [a for a in available_agents if chunk.required_capability in a.capabilities]
            if not compatible_agents:
                self.needs_more_agents = True
                continue
                
            agent = compatible_agents.pop(0)
            available_agents.remove(agent)

            agent.assign_chunk(chunk.chunk_id)
            agent.status = 'working'

            if chunk.depends_on:
                self.gated_agents.append(agent)
            else:
                agent.ready_to_execute = True

        if not verbose:
            self.print_summary()

        task = task_objects.get(self.task_id)
        if task and task.is_complete():
            task.mark_complete()
            self.should_shutdown = True
           
    def sync(self):
        pass

    def print_summary(self):
        task = get_task_by_id(self.task_id)
        total_chunks = len([c for c in chunk_objects.values() if c.task_id == self.task_id])
        pending_chunks = len([c for c in chunk_objects.values() if c.task_id == self.task_id and c.status == "pending"])
        completed_chunks = count_completed(self.task_id)
        summary = render_zone_summary(
            self.zone_id,
            task.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)