from hiveos.runtime.zone_controller import ZoneController
from hiveos.runtime.states import AgentStatus, ChunkStatus, TaskStatus, ZoneStallState
from threading import Event
from collections import deque
import time
import yaml

from hiveos.runtime.agent_runtime import flush_agent_db_state

from hiveos.runtime.chunk_runtime import flush_chunk_db_state
from hiveos.runtime.task_runtime import flush_task_db_state

from hiveos.messaging.message_bus import MessageBus
from hiveos.intelligence.execution_plane import AIExecutionPlane
from hiveos.intelligence.tool_registry import ToolCapabilityRegistry
from hiveos.intelligence.proposals import ProposalManager
from hiveos.intelligence.memory_store import MemoryStore
from hiveos.intelligence.sandbox_runtime import SandboxRuntimeManager
from hiveos.intelligence.tracing import build_trace_event, emit_trace
from hiveos.services.core_tool_service import CoreToolService
from hiveos.services.intent_service import IntentOrchestrationService
from hiveos.services.terminal_service import TerminalOrchestrationService
from hiveos.services.task_ingestion_service import TaskIngestionService
from hiveos.config import (
    ALLOW_WRAPPER_PY,
    AI_EXEC_MAX_QUEUE,
    AI_EXEC_MAX_RUNS,
    AI_EXEC_MAX_RUNTIME_SEC,
    AI_EXEC_QUEUE_TTL_SEC,
    REQUIRE_WRAPPER_SPEC_SIGNATURE,
    TICK_BUDGET_MS,
    TICK_METRICS_WINDOW,
    WRAPPER_SPEC_ALLOWLIST,
    WRAPPER_SPEC_SIGNING_KEY,
    ZONE_STALL_FAIL_TICKS,
    ZONE_STALL_WARN_TICKS,
)

class Core:
    def __init__(self, session_id=None):
        self.session_id = session_id
        self.chunk_objects = {}
        self.task_objects = {}
        self.task_name_map = {}  # task_name → list of task_ids
        self._runtime_agents = {}
        self.zones = {}
        self.zone_counter = 1
        self.global_bus = MessageBus()
        self.ai_execution_plane = AIExecutionPlane(
            max_runs=AI_EXEC_MAX_RUNS,
            max_queue=AI_EXEC_MAX_QUEUE,
            queue_ttl_seconds=AI_EXEC_QUEUE_TTL_SEC,
            max_runtime_seconds=AI_EXEC_MAX_RUNTIME_SEC,
        )
        self.ai_execution_plane.start()
        self.tool_registry = ToolCapabilityRegistry()
        self.proposal_manager = ProposalManager()
        self.memory_store = MemoryStore(session_id=self.session_id)
        self.sandbox_runtime = SandboxRuntimeManager()
        self.terminal_service = TerminalOrchestrationService(self)
        self.intent_service = IntentOrchestrationService(self)
        self.tool_service = CoreToolService(self)
        self.task_ingestion_service = TaskIngestionService(self)
        self.tool_service.register_builtin_tools()
        self.tick_control = {
            "running": True,
            "shutdown": False,  # ✅ Flag to signal shutdown
            "delay": 0.1,
            "stop_event": Event()
        }
        self.execution_hardening = {
            "stall_warn_ticks": max(1, int(ZONE_STALL_WARN_TICKS)),
            "stall_fail_ticks": max(0, int(ZONE_STALL_FAIL_TICKS)),
        }
        self.tick_counter = {"count": 0}
        self.tick_perf = {
            "window_size": max(20, int(TICK_METRICS_WINDOW)),
            "budget_ms": max(1, int(TICK_BUDGET_MS)),
            "tick_ms_window": deque(maxlen=max(20, int(TICK_METRICS_WINDOW))),
            "zone_ms_window": deque(maxlen=max(20, int(TICK_METRICS_WINDOW))),
            "flush_ms_window": deque(maxlen=max(20, int(TICK_METRICS_WINDOW))),
            "sandbox_ms_window": deque(maxlen=max(20, int(TICK_METRICS_WINDOW))),
            "over_budget_total": 0,
            "last_tick_ms": 0.0,
            "last_zone_ms": 0.0,
            "last_flush_ms": 0.0,
            "last_sandbox_ms": 0.0,
        }
        self.start_time = time.time()        
        self.global_bus.subscribe("zone/+/task_injection", self.handle_task_injection)
        self.global_bus.subscribe("zone/+/agent_injection", self.handle_agent_injection)



    def ingest_intent(self, intent):
        task_id = self.create_task(intent)
        self.delegate_task_to_zone(task_id)
        return task_id

    def create_task(self, intent):
        return self.task_ingestion_service.create_task(intent)

    def delegate_task_to_zone(self, task_id):
        zone_id = f"zone-{self.zone_counter}"
        self.zone_counter += 1
        zone = ZoneController(zone_id, self, global_bus=self.global_bus, session_id=self.session_id)
        zone.assign_task(task_id)
        

        needed = zone.agents_needed()
        agent_instances = [a.obj for a in self._runtime_agents.values() if a.obj.status == AgentStatus.IDLE]

        def is_already_assigned(agent_id):
            return any(agent_id in z.agent_ids for z in self.zones.values())

        filtered = [a for a in agent_instances if not is_already_assigned(a.agent_id)]

        required_caps = set(
            chunk.required_capability
            for chunk in self.chunk_objects.values()
            if chunk.task_id == task_id and chunk.status == ChunkStatus.PENDING
        )
        
        # Accept agents that match ANY required capability
        capable = [a for a in filtered if any(cap in a.capabilities for cap in required_caps)]
        
        if not capable:
            print(f"[Core] No capable agents found for zone {zone.zone_id} on task creation.")
        else:
            zone.assign_agents(capable[:needed], append=False)

        self.zones[zone_id] = zone
        task = self.task_objects.get(task_id)
        if task:
            task.zone_id = zone.zone_id

    def handle_task_injection(self, topic, task_reference):
        print(f"[Core] Received task injection with type: {type(task_reference)} value: {task_reference}")
        try:
            task_ids = self.task_name_map.get(task_reference)
            if not task_ids:
                print(f"[Core] No task found with name '{task_reference}'")
                return

            task_id = task_ids[-1]  # Use most recently injected task with this name
            task_obj = self.task_objects.get(task_id)
            if not task_obj:
                print(f"[Core] Task {task_id} not found in task_objects")
                return

            intent = task_obj.to_yaml_dict(self.chunk_objects)
            task_id = self.ingest_intent(intent)
            print(f"[Core] Task {task_id} reinjected from reference '{task_reference}' via zone: {topic}")
        except Exception as e:
            print(f"[Core] Failed to inject task from topic {topic}: {e}")

    def handle_agent_injection(self, topic, payload):
        print(f"[Core] Received agent injection request via {topic}")
        
        try:
            from hiveos.utilities.wrapper_serialization import (
                hydrate_wrapper,
                validate_py_wrapper,
                validate_wrapper_spec_policy,
            )

            if isinstance(payload, dict) and "wrapper_spec" in payload:
                spec = payload["wrapper_spec"]
                ok, reason = validate_wrapper_spec_policy(
                    spec,
                    allowlist=WRAPPER_SPEC_ALLOWLIST,
                    require_signature=REQUIRE_WRAPPER_SPEC_SIGNATURE,
                    signing_key=WRAPPER_SPEC_SIGNING_KEY,
                )
                if not ok:
                    print(f"[Core] Wrapper spec rejected: {reason}")
                    return
                agent = hydrate_wrapper(spec)

            elif isinstance(payload, dict) and "wrapper_py" in payload:
                if not ALLOW_WRAPPER_PY:
                    print("[Core] wrapper_py rejected: disabled by config")
                    return
                success, result = validate_py_wrapper(payload["wrapper_py"])
                if not success:
                    print(f"[Core] Wrapper validation failed: {result}")
                    return
                ok, reason = validate_wrapper_spec_policy(
                    result,
                    allowlist=WRAPPER_SPEC_ALLOWLIST,
                    require_signature=REQUIRE_WRAPPER_SPEC_SIGNATURE,
                    signing_key=WRAPPER_SPEC_SIGNING_KEY,
                )
                if not ok:
                    print(f"[Core] Wrapper from wrapper_py rejected: {reason}")
                    return
                agent = hydrate_wrapper(result)

            else:
                print(f"[Core] Invalid agent injection payload format: {payload}")
                return

            self.inject_agent(agent)
            print(f"[Core] Injected agent '{agent.agent_id}' with capabilities: {agent.capabilities}")

        except Exception as e:
            print(f"[Core] Agent injection failed: {e}")


    def tick(self, verbose=False):
        tick_start = time.perf_counter()
        assigned_ids = set()
        for z in self.zones.values():
            assigned_ids.update(z.agent_ids)

        # print(f"[Core] Zones in system: {list(self.zones.keys())}")
    

        for zone in list(self.zones.values()):

            # print(f"[Core] Ticking zone {zone.zone_id}. Frozen: {getattr(zone, 'frozen', False)}")
            
            if getattr(zone, "frozen", False):
                continue
                
            if zone.needs_more_agents:
                needed = zone.agents_needed()
                # print(f"[Debug] Zone {zone.zone_id} needs {needed} agents")
                # print(f"[Debug] Zone has agents: {[a.agent_id for a in zone.agent_instances]}")

                agent_instances = [a for a in self._runtime_agents.values() if a.obj.status == AgentStatus.IDLE]
                
                filtered = [a for a in agent_instances if a.agent_id not in assigned_ids]

                
                # Recompute required capabilities per zone
                #chunk_rows = Chunk.get_pending_by_task(zone.task_id)
                
                required_caps = set(
                    chunk.required_capability
                    for chunk in self.chunk_objects.values()
                    if chunk.task_id == zone.task_id and chunk.status == ChunkStatus.PENDING
                )
                
                capable = [a.obj for a in filtered if any(cap in a.obj.capabilities for cap in required_caps)]
                
                if capable:
                    zone.assign_agents(capable[:needed], append=True)
                    zone.needs_more_agents = False
                # else:
                    # print(f"[Core] No capable agents found for zone {zone.zone_id} on task creation.")
                    # Leave needs_more_agents = True so tick() can retry


            zone_report = zone.tick(verbose=verbose, current_tick=self.tick_counter["count"] + 1)
            self._emit_zone_watchdog_events(zone, zone_report)

            for agent in zone.releasable_agents():
                zone.release_agent(agent.agent_id)
                agent.status = AgentStatus.IDLE
                agent.zone = None

            if zone.should_shutdown:
                self.cleanup_zone(zone)
                
        
        zone_ms = (time.perf_counter() - tick_start) * 1000.0
        flush_start = time.perf_counter()
        flush_agent_db_state()
        flush_chunk_db_state(self.chunk_objects)
        flush_task_db_state(self.task_objects)
        flush_ms = (time.perf_counter() - flush_start) * 1000.0
        sandbox_start = time.perf_counter()
        self.sandbox_runtime.process_next()
        sandbox_ms = (time.perf_counter() - sandbox_start) * 1000.0
        self._emit_sandbox_runtime_events()
        tick_ms = (time.perf_counter() - tick_start) * 1000.0
        self._record_tick_perf(tick_ms=tick_ms, zone_ms=zone_ms, flush_ms=flush_ms, sandbox_ms=sandbox_ms)

    def cleanup_zone(self, zone):
        # Return agents to pool
        for agent in zone.agent_instances:
            agent.status = AgentStatus.IDLE
            agent.zone = None
            agent.current_chunk = None
            
        # Set task to terminated if still running
        task = self.task_objects.get(zone.task_id)
        if task and task.status not in (TaskStatus.COMPLETED, TaskStatus.TERMINATED, TaskStatus.FAILED):
            task.status = TaskStatus.TERMINATED
            
        # Remove zone    
        del self.zones[zone.zone_id]
        print(f"[Core] Zone {zone.zone_id} has completed and been cleaned up.")
        
    def remove_agent(self, agent_id):
        runtime_agent = self._runtime_agents.pop(agent_id, None)
        agent = runtime_agent.obj if runtime_agent else None

        if not agent:
            return False

        for zone in self.zones.values():
            if agent_id in zone.agent_ids:
                zone.release_agent(agent_id)
                zone.needs_more_agents = True

                # Unassign chunks that were assigned to this agent
        for chunk in self.chunk_objects.values():
            if chunk.assigned_agent == agent_id:
                if chunk.status in (ChunkStatus.ASSIGNED, ChunkStatus.PENDING):
                    chunk.assigned_agent = None
                    chunk.status = ChunkStatus.PENDING
                    print(f"[Core] Chunk {chunk.chunk_id} released due to agent {agent_id} removal.")
                
        print(f"[Core] Agent {agent_id} removed from runtime and detached from zones.")        
        return True
        
    def remove_task(self, task_id):
        # Find and cleanup the zone if the task owns one
        for zone in list(self.zones.values()):
            if zone.task_id == task_id:
                self.cleanup_zone(zone)

        # Remove task object
        task = self.task_objects.pop(task_id, None)
        if task:
            print(f"[Core] Task {task_id} removed from task_objects")

        # Remove chunk objects associated with task
        to_delete = [cid for cid, c in self.chunk_objects.items() if c.task_id == task_id]
        for cid in to_delete:
            del self.chunk_objects[cid]
            print(f"[Core] Chunk {cid} removed from chunk_objects")

    def inject_agent(self, agent_obj):
        from hiveos.runtime.runtime_agent import register
        if hasattr(agent_obj, "bind_core"):
            agent_obj.bind_core(self)
        if hasattr(agent_obj, "bind_execution_plane"):
            agent_obj.bind_execution_plane(self.ai_execution_plane)
        if hasattr(agent_obj, "bind_tool_registry"):
            agent_obj.bind_tool_registry(self.tool_registry)
        if hasattr(agent_obj, "bind_proposal_manager"):
            agent_obj.bind_proposal_manager(self.proposal_manager, self)
        register(agent_obj, self)

    def _emit_zone_watchdog_events(self, zone, zone_report):
        if not isinstance(zone_report, dict) or not zone_report.get("state_changed"):
            return
        state = str(zone_report.get("stall_state") or "")
        event_type = None
        outcome = "success"
        if state == ZoneStallState.BLOCKED:
            event_type = "zone_execution_blocked"
            outcome = "denied"
        elif state == ZoneStallState.HEALTHY:
            event_type = "zone_execution_recovered"
            outcome = "success"
        elif state == ZoneStallState.FAILED:
            event_type = "zone_execution_failed"
            outcome = "failed"
        if not event_type:
            return
        emit_trace(
            build_trace_event(
                event_type=event_type,
                run_id=f"{zone.zone_id}:watchdog",
                outcome=outcome,
                payload={
                    "zone_id": zone.zone_id,
                    "task_id": zone.task_id,
                    "stall_state": state,
                    "stall_ticks": zone_report.get("stall_ticks"),
                    "blocked_reasons": zone_report.get("blocked_reasons"),
                    "blocked_reason_counts": zone_report.get("blocked_reason_counts"),
                    "warn_ticks": zone.stall_warn_ticks,
                    "fail_ticks": zone.stall_fail_ticks,
                },
                task_id=zone.task_id,
                zone_id=zone.zone_id,
            )
        )

    def get_runtime_perf(self):
        perf = self.tick_perf
        tick_p95 = _percentile(perf["tick_ms_window"], 95)
        budget = float(perf["budget_ms"])
        return {
            "budget_ms": perf["budget_ms"],
            "window_size": perf["window_size"],
            "samples": len(perf["tick_ms_window"]),
            "last_tick_ms": round(float(perf["last_tick_ms"]), 3),
            "last_zone_ms": round(float(perf["last_zone_ms"]), 3),
            "last_flush_ms": round(float(perf["last_flush_ms"]), 3),
            "last_sandbox_ms": round(float(perf["last_sandbox_ms"]), 3),
            "tick_avg_ms": round(_avg(perf["tick_ms_window"]), 3),
            "tick_p50_ms": round(_percentile(perf["tick_ms_window"], 50), 3),
            "tick_p95_ms": round(tick_p95, 3),
            "zone_avg_ms": round(_avg(perf["zone_ms_window"]), 3),
            "flush_avg_ms": round(_avg(perf["flush_ms_window"]), 3),
            "sandbox_avg_ms": round(_avg(perf["sandbox_ms_window"]), 3),
            "over_budget_total": int(perf["over_budget_total"]),
            "over_budget_ratio": round(_safe_div(perf["over_budget_total"], self.tick_counter["count"]), 4),
            "currently_over_budget": tick_p95 > budget if perf["tick_ms_window"] else False,
        }

    def _record_tick_perf(self, tick_ms, zone_ms, flush_ms, sandbox_ms):
        perf = self.tick_perf
        perf["last_tick_ms"] = tick_ms
        perf["last_zone_ms"] = zone_ms
        perf["last_flush_ms"] = flush_ms
        perf["last_sandbox_ms"] = sandbox_ms
        perf["tick_ms_window"].append(float(tick_ms))
        perf["zone_ms_window"].append(float(zone_ms))
        perf["flush_ms_window"].append(float(flush_ms))
        perf["sandbox_ms_window"].append(float(sandbox_ms))
        if tick_ms > perf["budget_ms"]:
            perf["over_budget_total"] += 1

    def _emit_sandbox_runtime_events(self):
        event_type_map = {
            "completed": "sandbox_run_completed",
            "failed": "sandbox_run_failed",
            "timeout": "sandbox_run_timeout",
            "isolation_applied": "sandbox_isolation_applied",
            "isolation_degraded": "sandbox_isolation_degraded",
        }
        for event in self.sandbox_runtime.drain_events(max_events=100):
            mapped = event_type_map.get(event.get("type"))
            if not mapped:
                continue
            emit_trace(
                build_trace_event(
                    event_type=mapped,
                    run_id=event.get("sandbox_run_id") or "sandbox",
                    outcome="success" if event.get("type") in {"started", "completed", "canceled"} else "failed",
                    payload={
                        "sandbox_run_id": event.get("sandbox_run_id"),
                        "sandbox_event_type": event.get("type"),
                        "timestamp": event.get("timestamp"),
                    },
                )
            )

    def start_background_tick(self):
        from threading import Thread

        def run():
            pause_logged = False
            while not self.tick_control["shutdown"]:  # ✅ Clean exit condition
                if self.tick_control["running"]:
                    pause_logged = False
                    try:
                        self.tick()
                        self.tick_counter["count"] += 1
                    except Exception as e:
                        print(f"[Tick Loop Error] {e}")
                        import traceback
                        traceback.print_exc()
                    self.tick_control["stop_event"].wait(self.tick_control["delay"])
                else:
                    if not pause_logged:
                        print(f"[Core {self.session_id}] Paused — waiting to resume")
                        pause_logged = True
                    self.tick_control["stop_event"].wait(0.1)  # ✅ Wait, don’t spin

            print(f"[Core {self.session_id}] Tick loop exited cleanly.")

        self.tick_thread = Thread(target=run, daemon=True)
        self.tick_thread.start()



    def teardown(self):
        print(f"[Core] Shutting down session {self.session_id}")

        self.tick_control["running"] = False
        self.tick_control["shutdown"] = True  # ✅ Allow tick thread to exit
        self.tick_control["stop_event"].set()  # ✅ Unblocks any current wait()

        # Cleanup all zones and agents
        for zone in list(self.zones.values()):
            self.cleanup_zone(zone)

        # Clear all runtime objects
        self.chunk_objects.clear()
        self.task_objects.clear()
        self._runtime_agents.clear()
        self.zones.clear()
        self.ai_execution_plane.stop()


def _avg(values):
    if not values:
        return 0.0
    return float(sum(values) / len(values))


def _percentile(values, pct):
    if not values:
        return 0.0
    ordered = sorted(float(v) for v in values)
    if len(ordered) == 1:
        return ordered[0]
    idx = int(round((pct / 100.0) * (len(ordered) - 1)))
    idx = max(0, min(len(ordered) - 1, idx))
    return ordered[idx]


def _safe_div(numerator, denominator):
    if not denominator:
        return 0.0
    return float(numerator) / float(denominator)
