from hiveos.messaging.message_bus import MessageBus
from hiveos.runtime.states import AgentStatus, ChunkStatus, TaskStatus
from hiveos.runtime.zone_watchdog import ZoneWatchdog
from hiveos.utilities.gate_evaluator import GateEvaluator
from hiveos.utilities.render_zone_summary import render_zone_summary


class ZoneController:
    def __init__(self, zone_id, core, global_bus=None, session_id=None):
        self.core = core
        self.zone_id = zone_id
        self.session_id = session_id
        self.task_id = None
        self.required_capability = "generic"
        self.agent_ids = []
        self.agent_instances = []
        self.chunk_results = {}
        self.should_shutdown = False
        self.needs_more_agents = False
        self.fallback_buffer = 2
        self.gated_agents = []
        self.local_bus = MessageBus()
        self.global_bus = global_bus or {}
        self.gate_evaluator = GateEvaluator(self)
        self.frozen = False
        self.pending_task_reference = None
        self._last_need_snapshot = None
        self._last_summary_snapshot = None

        hardening = getattr(core, "execution_hardening", {})
        self.watchdog = ZoneWatchdog(
            zone=self,
            stall_warn_ticks=hardening.get("stall_warn_ticks", 200),
            stall_fail_ticks=hardening.get("stall_fail_ticks", 0),
        )
        self.stall_warn_ticks = self.watchdog.stall_warn_ticks
        self.stall_fail_ticks = self.watchdog.stall_fail_ticks

        self.local_bus.subscribe(f"zone/{self.zone_id}/task_injection", self.capture_task_injection)
        self.local_bus.subscribe(f"zone/{self.zone_id}/agent_injection", self.capture_agent_injection)

    @property
    def stall_state(self):
        return self.watchdog.stall_state

    def assign_task(self, task_id):
        self.task_id = task_id
        if not self.agent_instances:
            self.needs_more_agents = True
            print(f"[Zone {self.zone_id}] No agents assigned yet, will need to wait for agents to be available.")

        self.chunks = {
            c.chunk_id: c
            for c in self.core.chunk_objects.values()
            if c.task_id == task_id and c.status == ChunkStatus.PENDING
        }

        first = next(iter(self.chunks.values()), None)
        self.required_capability = first.required_capability if first else "generic"

    def capture_task_injection(self, topic, task_reference):
        if self.pending_task_reference is not None:
            print(f"[Zone {self.zone_id}] Warning: multiple injections attempted - ignoring extra")
            return
        self.pending_task_reference = task_reference

    def capture_agent_injection(self, topic, wrapper_payload):
        print(f"[Zone {self.zone_id}] Received agent injection payload, forwarding to core.")
        self.global_bus.publish(f"zone/{self.zone_id}/agent_injection", wrapper_payload)

    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.message_bus = self.local_bus

    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 = [existing for existing in self.agent_ids if existing != agent_id]

    def agents_needed(self):
        pending_chunks = [
            c for c in self.core.chunk_objects.values()
            if c.task_id == self.task_id and c.status == ChunkStatus.PENDING
        ]
        idle_agents = [a for a in self.agent_instances if a.status == AgentStatus.IDLE]
        need = max(len(pending_chunks) - len(idle_agents), 0)
        snapshot = (len(pending_chunks), len(idle_agents), need)
        if snapshot != self._last_need_snapshot:
            print(
                f"[Zone {self.zone_id}] pending_chunks: {len(pending_chunks)}, "
                f"idle_agents: {len(idle_agents)}, need: {need}"
            )
            self._last_need_snapshot = snapshot
        return need + self.fallback_buffer

    def releasable_agents(self):
        pending_chunks = [
            c for c in self.core.chunk_objects.values()
            if c.task_id == self.task_id and c.status == ChunkStatus.PENDING
        ]
        idle_agents = [a for a in self.agent_instances if a.status == AgentStatus.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, current_tick=None):
        task = self.core.task_objects.get(self.task_id)
        if not task:
            return self.watchdog.build_tick_report(made_progress=False, state_changed=False)

        made_progress = False
        before = self.watchdog.build_progress_fingerprint(task)
        pending = self._pending_chunks()
        made_progress = self.watchdog.propagate_stale_dependencies(pending) or made_progress
        pending = self._pending_chunks()

        for agent in self.agent_instances:
            if agent.status == AgentStatus.WORKING and not agent.ready_to_execute:
                chunk = self.core.chunk_objects.get(agent.current_chunk)
                if not chunk:
                    print(f"[GateEval] Agent {agent.agent_id} assigned to unknown chunk ID {agent.current_chunk}")
                    continue
                if self.gate_evaluator.should_execute(chunk):
                    agent.ready_to_execute = True

        for agent in self.agent_instances:
            if agent.status == AgentStatus.WORKING and agent.ready_to_execute:
                chunk = self.core.chunk_objects.get(agent.current_chunk)
                if chunk and chunk.depends_on:
                    dep = self.core.chunk_objects.get(chunk.depends_on)
                    if not dep or dep.status != ChunkStatus.COMPLETED:
                        continue
                agent.tick()

        available_agents = [a for a in self.agent_instances if a.status == AgentStatus.IDLE and a.current_chunk is None]
        for chunk in pending:
            if task.has_conditionals and not self.gate_evaluator.should_assign(chunk):
                continue

            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)
            if task.status == TaskStatus.PENDING:
                task.status = TaskStatus.RUNNING
            agent.assign_chunk(chunk)
            agent.status = AgentStatus.WORKING
            agent.ready_to_execute = self.gate_evaluator.should_execute(chunk)

        after = self.watchdog.build_progress_fingerprint(task)
        if before != after:
            made_progress = True

        pending = self._pending_chunks()
        self.watchdog.refresh_blocked_reasons(task, pending)
        state_changed = self.watchdog.update_stall_state(
            task=task,
            pending_chunks=pending,
            made_progress=made_progress,
            current_tick=current_tick,
        )

        if not verbose:
            self.print_summary()

        if task.is_complete(self.core.chunk_objects):
            if task.status not in (TaskStatus.FAILED, TaskStatus.TERMINATED):
                task.mark_complete()
            self.should_shutdown = True
            if self.pending_task_reference:
                self.global_bus.publish(f"zone/{self.zone_id}/task_injection", self.pending_task_reference)

        return self.watchdog.build_tick_report(made_progress=made_progress, state_changed=state_changed)

    def print_summary(self):
        task = self.core.task_objects.get(self.task_id)
        if not task:
            return
        total_chunks = len([c for c in self.core.chunk_objects.values() if c.task_id == self.task_id])
        completed_chunks = sum(
            c.status == ChunkStatus.COMPLETED
            for c in self.core.chunk_objects.values()
            if c.task_id == self.task_id
        )
        summary_snapshot = (
            total_chunks,
            completed_chunks,
            len(self.agent_instances),
            sum(a.status == AgentStatus.WORKING for a in self.agent_instances),
            sum(a.status == AgentStatus.IDLE for a in self.agent_instances),
            sum(a.status == AgentStatus.RECHARGING for a in self.agent_instances),
            sum(a.status == AgentStatus.COOLDOWN for a in self.agent_instances),
        )
        if summary_snapshot == self._last_summary_snapshot:
            return
        self._last_summary_snapshot = summary_snapshot
        summary = render_zone_summary(
            self.zone_id,
            task.name,
            total_chunks,
            completed_chunks,
            {
                "total": len(self.agent_instances),
                "working": sum(a.status == AgentStatus.WORKING for a in self.agent_instances),
                "idle": sum(a.status == AgentStatus.IDLE for a in self.agent_instances),
                "recharging": sum(a.status == AgentStatus.RECHARGING for a in self.agent_instances),
                "cooldown": sum(a.status == AgentStatus.COOLDOWN for a in self.agent_instances),
            },
        )
        try:
            print(summary)
        except UnicodeEncodeError:
            # Some Windows terminals default to cp1252 and cannot print all glyphs.
            print(summary.encode("ascii", errors="replace").decode("ascii"))

    def get_watchdog_summary(self):
        return self.watchdog.get_summary()

    def get_chunk(self, chunk_id):
        return self.core.chunk_objects.get(chunk_id)

    def get_chunk_result(self, chunk_id):
        return self.chunk_results.get(chunk_id)

    def _pending_chunks(self):
        return [
            c for c in self.core.chunk_objects.values()
            if c.task_id == self.task_id and c.status == ChunkStatus.PENDING
        ]
