from hiveos.sim_agent import SimAgent
from hiveos.db import get_db

# current
agent_objects = {}

# future
zone_agent_pools = {
    "zone1": {},
    "zone2": {},
}

def register(agent):
    agent_objects[agent.agent_id] = agent
    try:
        with get_db() as conn:
            conn.execute("""
                INSERT OR REPLACE INTO agents (agent_id, status, battery, zone, capabilities, zone_residency_until)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                agent.agent_id,
                agent.status,
                100,  # default battery
                agent.zone,
                str(agent.capabilities),
                None
            ))
    except Exception as e:
        print(f"[WARN] Failed to register agent in DB: {e}")

def flush_agent_db_state():
    
    try:
        with get_db() as conn:
            for agent in agent_objects.values():
                conn.execute("""
                    UPDATE agents
                    SET status = ?, zone = ?, capabilities = ?
                    WHERE agent_id = ?
                """, (
                    agent.status,
                    agent.zone,
                    str(agent.capabilities),
                    agent.agent_id
                ))
    except Exception as e:
        print(f"[WARN] Failed to flush agent DB state: {e}")

    

def get_by_id(agent_id):
    return agent_objects.get(agent_id)
    
def get_idle_agents(n=5):
    return [
        a for a in agent_objects.values()
        if a.status == "idle" and a.zone is None
    ][:n]
    
def deregister_agent(agent_id):
    agent = agent_objects.pop(agent_id, None)
    if not agent:
        print(f"[Runtime] Agent '{agent_id}' not found.")
        return None

    # Reset any chunk it was working on
    from hiveos.runtime.chunk_runtime import reset, get_task_id
    from hiveos.utilities.log_event import log_chunk_event
    if agent.current_chunk:
        task_id = get_task_id(agent.current_chunk)
        reset(agent.current_chunk)
        log_chunk_event(agent.current_chunk, task_id, agent_id, "fallback", reason="manual removal")
        
    try:        
        with get_db() as conn:
            conn.execute("DELETE FROM agents WHERE agent_id = ?", (agent_id,))
    except Exception as e:
        print(f"[WARN] Failed to delete agent '{agent_id}' from DB: {e}")

    print(f"[Runtime] Agent '{agent_id}' deregistered and removed from database.")
    return agent
