-- Function to handle atomic node heartbeat and metrics update
-- Updated to include Swapper metrics
CREATE OR REPLACE FUNCTION fn_node_heartbeat(
    p_node_id text,
    p_metrics jsonb
) RETURNS void AS $$
BEGIN
    INSERT INTO node_cache (
        node_id, node_name, hostname, roles, environment, status, gpu, last_sync,
        cpu_model, cpu_cores, cpu_usage, gpu_model, gpu_vram_total, gpu_vram_used,
        ram_total, ram_used, disk_total, disk_used, agent_count_router, agent_count_system,
        last_heartbeat, dagi_router_url, self_healing_status,
        swapper_healthy, swapper_models_loaded, swapper_models_total
    )
    VALUES (
        p_node_id,
        p_metrics->>'node_name',
        p_metrics->>'hostname',
        (SELECT array_agg(elem) FROM jsonb_array_elements_text(p_metrics->'roles') AS elem),
        p_metrics->>'environment',
        p_metrics->>'status',
        p_metrics->>'gpu',
        NOW(),
        p_metrics->>'cpu_model',
        (p_metrics->>'cpu_cores')::integer,
        (p_metrics->>'cpu_usage')::numeric,
        p_metrics->>'gpu_model',
        (p_metrics->>'gpu_vram_total')::integer,
        (p_metrics->>'gpu_vram_used')::integer,
        (p_metrics->>'ram_total')::integer,
        (p_metrics->>'ram_used')::integer,
        (p_metrics->>'disk_total')::integer,
        (p_metrics->>'disk_used')::integer,
        (p_metrics->>'agent_count_router')::integer,
        (p_metrics->>'agent_count_system')::integer,
        NOW(),
        p_metrics->>'dagi_router_url',
        p_metrics->>'self_healing_status',
        (p_metrics->>'swapper_healthy')::boolean,
        (p_metrics->>'swapper_models_loaded')::integer,
        (p_metrics->>'swapper_models_total')::integer
    )
    ON CONFLICT (node_id) DO UPDATE SET
        node_name = COALESCE(EXCLUDED.node_name, node_cache.node_name),
        hostname = COALESCE(EXCLUDED.hostname, node_cache.hostname),
        roles = COALESCE(EXCLUDED.roles, node_cache.roles),
        environment = COALESCE(EXCLUDED.environment, node_cache.environment),
        status = COALESCE(EXCLUDED.status, node_cache.status),
        gpu = COALESCE(EXCLUDED.gpu, node_cache.gpu),
        last_sync = NOW(),
        cpu_model = COALESCE(EXCLUDED.cpu_model, node_cache.cpu_model),
        cpu_cores = COALESCE(EXCLUDED.cpu_cores, node_cache.cpu_cores),
        cpu_usage = COALESCE(EXCLUDED.cpu_usage, node_cache.cpu_usage),
        gpu_model = COALESCE(EXCLUDED.gpu_model, node_cache.gpu_model),
        gpu_vram_total = COALESCE(EXCLUDED.gpu_vram_total, node_cache.gpu_vram_total),
        gpu_vram_used = COALESCE(EXCLUDED.gpu_vram_used, node_cache.gpu_vram_used),
        ram_total = COALESCE(EXCLUDED.ram_total, node_cache.ram_total),
        ram_used = COALESCE(EXCLUDED.ram_used, node_cache.ram_used),
        disk_total = COALESCE(EXCLUDED.disk_total, node_cache.disk_total),
        disk_used = COALESCE(EXCLUDED.disk_used, node_cache.disk_used),
        agent_count_router = COALESCE(EXCLUDED.agent_count_router, node_cache.agent_count_router),
        agent_count_system = COALESCE(EXCLUDED.agent_count_system, node_cache.agent_count_system),
        last_heartbeat = NOW(),
        dagi_router_url = COALESCE(EXCLUDED.dagi_router_url, node_cache.dagi_router_url),
        self_healing_status = COALESCE(EXCLUDED.self_healing_status, node_cache.self_healing_status),
        swapper_healthy = COALESCE(EXCLUDED.swapper_healthy, node_cache.swapper_healthy),
        swapper_models_loaded = COALESCE(EXCLUDED.swapper_models_loaded, node_cache.swapper_models_loaded),
        swapper_models_total = COALESCE(EXCLUDED.swapper_models_total, node_cache.swapper_models_total),
        updated_at = NOW();
END;
$$ LANGUAGE plpgsql;