from collections import Counter

from hiveos.runtime.states import AgentStatus, ChunkStatus, TaskStatus, ZoneStallState


class ZoneWatchdog:
    def __init__(self, zone, stall_warn_ticks=200, stall_fail_ticks=0):
        self.zone = zone
        self.stall_warn_ticks = max(1, int(stall_warn_ticks))
        self.stall_fail_ticks = max(0, int(stall_fail_ticks))
        self.stall_ticks = 0
        self.blocked_since_tick = None
        self.blocked_reasons = []
        self.blocked_reason_counts = {}
        self.stall_state = ZoneStallState.HEALTHY

    def build_progress_fingerprint(self, task):
        chunks = self._task_chunks()
        status_counter = Counter(chunk.status for chunk in chunks)
        work = []
        for agent in self.zone.agent_instances:
            work.append(
                (
                    agent.agent_id,
                    agent.status,
                    agent.current_chunk,
                    bool(getattr(agent, "ready_to_execute", False)),
                    bool(getattr(agent, "waiting_on_input", False)),
                )
            )
        work.sort()
        return (
            task.status,
            tuple(sorted(status_counter.items())),
            tuple(work),
        )

    def propagate_stale_dependencies(self, pending_chunks):
        changed = False
        for chunk in pending_chunks:
            dep_id = chunk.depends_on
            if not dep_id:
                continue
            dep_chunk = self.zone.core.chunk_objects.get(dep_id)
            if dep_chunk and dep_chunk.status == ChunkStatus.STALE:
                chunk.mark_stale()
                changed = True
        return changed

    def refresh_blocked_reasons(self, task, pending_chunks):
        reasons = []
        for chunk in pending_chunks:
            reason = self._classify_pending_chunk_reason(task, chunk)
            reasons.append(reason)

        for agent in self.zone.agent_instances:
            if agent.status != AgentStatus.WORKING:
                continue
            if bool(getattr(agent, "waiting_on_input", False)):
                reasons.append("waiting_input")
            elif not bool(getattr(agent, "ready_to_execute", False)):
                reasons.append("execution_gate")

        if not reasons and pending_chunks:
            reasons.append("unknown_pending")

        self.blocked_reasons = reasons
        self.blocked_reason_counts = dict(Counter(reasons))

    def update_stall_state(self, task, pending_chunks, made_progress, current_tick):
        previous_state = self.stall_state

        if not pending_chunks:
            self.stall_ticks = 0
            self.blocked_since_tick = None
            self.stall_state = ZoneStallState.HEALTHY
            if task.status == TaskStatus.BLOCKED:
                task.status = TaskStatus.RUNNING
            return previous_state != self.stall_state

        if made_progress:
            self.stall_ticks = 0
            self.blocked_since_tick = None
            if task.status == TaskStatus.BLOCKED:
                task.status = TaskStatus.RUNNING
            self.stall_state = ZoneStallState.HEALTHY
            return previous_state != self.stall_state

        self.stall_ticks += 1
        if self.blocked_since_tick is None and current_tick is not None:
            self.blocked_since_tick = int(current_tick)

        if self.stall_ticks >= self.stall_warn_ticks:
            self.stall_state = ZoneStallState.BLOCKED
            if task.status in (TaskStatus.PENDING, TaskStatus.RUNNING):
                task.status = TaskStatus.BLOCKED
        else:
            self.stall_state = ZoneStallState.HEALTHY

        if self.stall_fail_ticks > 0 and self.stall_ticks >= self.stall_fail_ticks:
            for chunk in pending_chunks:
                chunk.mark_stale()
            task.status = TaskStatus.FAILED
            self.stall_state = ZoneStallState.FAILED
            self.zone.should_shutdown = True

        return previous_state != self.stall_state

    def build_tick_report(self, made_progress, state_changed):
        return {
            "zone_id": self.zone.zone_id,
            "task_id": self.zone.task_id,
            "made_progress": bool(made_progress),
            "stall_state": self.stall_state,
            "stall_ticks": int(self.stall_ticks),
            "state_changed": bool(state_changed),
            "blocked_reasons": list(self.blocked_reasons),
            "blocked_reason_counts": dict(self.blocked_reason_counts),
        }

    def get_summary(self):
        return {
            "state": self.stall_state,
            "stall_ticks": int(self.stall_ticks),
            "warn_ticks": int(self.stall_warn_ticks),
            "fail_ticks": int(self.stall_fail_ticks),
            "blocked_since_tick": self.blocked_since_tick,
            "blocked_reasons": list(self.blocked_reasons),
            "blocked_reason_counts": dict(self.blocked_reason_counts),
        }

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

    def _task_chunks(self):
        return [c for c in self.zone.core.chunk_objects.values() if c.task_id == self.zone.task_id]

    def _classify_pending_chunk_reason(self, task, chunk):
        dep_id = chunk.depends_on
        if dep_id:
            dep_chunk = self.zone.core.chunk_objects.get(dep_id)
            if dep_chunk is None:
                return "dependency_missing"
            if dep_chunk.status == ChunkStatus.STALE:
                return "dependency_stale"
            if dep_chunk.status != ChunkStatus.COMPLETED:
                return "dependency_pending"

        cond = chunk.conditional
        if task.has_conditionals and isinstance(cond, dict) and cond.get("type"):
            key = str(cond.get("key") or "").strip()
            if key and self.zone.local_bus.get(key) is None:
                return "conditional_unresolved"
            return "conditional_blocked"

        cap = str(chunk.required_capability or "generic")
        if not self._capability_exists_in_session(cap):
            return f"missing_capability:{cap}"
        if not self._idle_capability_available_in_zone(cap):
            return f"waiting_capability_agent:{cap}"
        return "assignable_pending"

    def _capability_exists_in_session(self, capability):
        for runtime_agent in self.zone.core._runtime_agents.values():
            if runtime_agent.session_id != self.zone.session_id:
                continue
            if capability in list(getattr(runtime_agent.obj, "capabilities", []) or []):
                return True
        return False

    def _idle_capability_available_in_zone(self, capability):
        return any(
            agent.status == AgentStatus.IDLE and capability in list(getattr(agent, "capabilities", []) or [])
            for agent in self.zone.agent_instances
        )
