app/core/mesh.py

import time
from typing import Dict, List, Optional, Set
from dataclasses import dataclass, field
from enum import Enum

from app.core.models import JobType, WorkerRuntimeType, NodeType
from app.phone.workers import get_worker, list_workers_for_session
from app.phone.jobs import create_job, get_job, complete_job
from app.gbridge.leases import list_active_leases


class MeshNodeStatus(str, Enum):
    ONLINE = "online"
    BUSY = "busy"
    OFFLINE = "offline"


@dataclass
class MeshNode:
    node_id: str
    node_type: NodeType
    runtime_type: WorkerRuntimeType
    capabilities: List[str] = field(default_factory=list)
    status: MeshNodeStatus = MeshNodeStatus.ONLINE
    last_heartbeat: float = field(default_factory=time.time)
    tokens_per_second: float = 0.0
    ram_used_mb: float = 0.0
    ram_total_mb: float = 0.0
    cpu_usage: float = 0.0
    thermal_state: str = "nominal"
    battery_level: float = 100.0
    current_job_id: Optional[str] = None
    supported_models: List[str] = field(default_factory=list)
    max_tokens_per_job: int = 512


class MeshRegistry:
    """Global mesh registry for all compute nodes across all sessions."""

    _nodes: Dict[str, MeshNode] = {}
    _session_nodes: Dict[str, Set[str]] = {}

    @classmethod
    def register_node(cls, node: MeshNode) -> None:
        cls._nodes[node.node_id] = node

    @classmethod
    def update_node_heartbeat(cls, node_id: str, metrics: dict) -> None:
        node = cls._nodes.get(node_id)
        if node:
            node.last_heartbeat = time.time()
            node.tokens_per_second = metrics.get("tokens_per_second", 0.0)
            node.ram_used_mb = metrics.get("ram_used_mb", 0.0)
            node.ram_total_mb = metrics.get("ram_total_mb", 0.0)
            node.cpu_usage = metrics.get("cpu_usage", 0.0)
            node.thermal_state = metrics.get("thermal_state", "nominal")
            node.battery_level = metrics.get("battery_level", 100.0)
            node.status = MeshNodeStatus.BUSY if metrics.get("is_generating") else MeshNodeStatus.ONLINE
            if node.current_job_id and not metrics.get("is_generating"):
                node.current_job_id = None

    @classmethod
    def mark_node_busy(cls, node_id: str, job_id: str) -> None:
        node = cls._nodes.get(node_id)
        if node:
            node.status = MeshNodeStatus.BUSY
            node.current_job_id = job_id

    @classmethod
    def mark_node_free(cls, node_id: str) -> None:
        node = cls._nodes.get(node_id)
        if node:
            node.status = MeshNodeStatus.ONLINE
            node.current_job_id = None

    @classmethod
    def get_node(cls, node_id: str) -> Optional[MeshNode]:
        return cls._nodes.get(node_id)

    @classmethod
    def list_nodes(cls) -> List[MeshNode]:
        cls._prune_stale_nodes()
        return list(cls._nodes.values())

    @classmethod
    def list_available_nodes(cls, capability: Optional[str] = None) -> List[MeshNode]:
        cls._prune_stale_nodes()
        nodes = [n for n in cls._nodes.values() if n.status != MeshNodeStatus.OFFLINE]
        if capability:
            nodes = [n for n in nodes if capability in n.capabilities]
        return nodes

    @classmethod
    def get_mesh_for_session(cls, session_id: str) -> List[MeshNode]:
        """Return all mesh nodes visible to a session (global mesh)."""
        cls._prune_stale_nodes()
        return list(cls._nodes.values())

    @classmethod
    def schedule_job(
        cls,
        job_type: JobType,
        session_id: str,
        payload: dict,
        strategy: str = "best_node"
    ) -> Optional[str]:
        """Schedule a job across the mesh and return assigned node_id."""
        cls._prune_stale_nodes()
        available = cls.list_available_nodes()
        if not available:
            return None

        if strategy == "best_node":
            # Pick node with highest tokens_per_second that supports this job type
            candidates = [n for n in available if cls._node_supports(n, job_type)]
            if not candidates:
                return None
            best = max(candidates, key=lambda n: n.tokens_per_second or 1.0)
            return best.node_id

        elif strategy == "round_robin":
            candidates = [n for n in available if cls._node_supports(n, job_type)]
            if not candidates:
                return None
            return candidates[0].node_id

        elif strategy == "parallel_split":
            # For split strategy, return the primary node (Mac/Ollama)
            # The secondary task (embedding/classification) goes to iPhone
            mac_nodes = [n for n in available if n.node_type == NodeType.MAC_AGENT and cls._node_supports(n, job_type)]
            if mac_nodes:
                return mac_nodes[0].node_id
            return available[0].node_id if available else None

        return None

    @classmethod
    def get_combined_throughput(cls) -> dict:
        """Return aggregate mesh stats."""
        nodes = cls.list_nodes()
        total_tps = sum(n.tokens_per_second for n in nodes)
        total_ram_used = sum(n.ram_used_mb for n in nodes)
        total_ram = sum(n.ram_total_mb for n in nodes)
        active_jobs = sum(1 for n in nodes if n.status == MeshNodeStatus.BUSY)
        return {
            "nodes_online": len([n for n in nodes if n.status != MeshNodeStatus.OFFLINE]),
            "nodes_busy": active_jobs,
            "total_nodes": len(nodes),
            "combined_tokens_per_second": round(total_tps, 1),
            "combined_ram_used_mb": round(total_ram_used, 0),
            "combined_ram_total_mb": round(total_ram, 0),
            "active_jobs": active_jobs,
        }

    @classmethod
    def _node_supports(cls, node: MeshNode, job_type: JobType) -> bool:
        cap_map = {
            JobType.TEXT_EMBEDDING: ["iphone.text.echo.private", "mac.embed.text.local"],
            JobType.SMALL_LLM_GENERATE: ["iphone.text.echo.private", "mac.llm.generate.local"],
            JobType.PRIVACY_REDACTION: ["iphone.privacy.redact.local"],
            JobType.IMAGE_CLASSIFICATION: ["iphone.image.classify.local"],
        }
        required_caps = cap_map.get(job_type, [])
        return any(cap in node.capabilities for cap in required_caps) or not required_caps

    @classmethod
    def _prune_stale_nodes(cls, timeout_seconds: float = 60.0) -> None:
        now = time.time()
        stale = [nid for nid, n in cls._nodes.items() if now - n.last_heartbeat > timeout_seconds]
        for nid in stale:
            cls._nodes[nid].status = MeshNodeStatus.OFFLINE


# Initialize mesh from existing workers
def sync_mesh_from_workers() -> None:
    """Sync mesh registry with existing phone worker state."""
    from app.phone.workers import _workers_by_session
    for session_id, workers in _workers_by_session.items():
        for worker in workers.values():
            node = MeshNode(
                node_id=worker.worker_id,
                node_type=NodeType.IPHONE_WORKER,
                runtime_type=worker.runtime_type or WorkerRuntimeType.SAFARI_WASM,
                capabilities=[c.capability_name for c in (worker.capabilities or [])],
                status=MeshNodeStatus.ONLINE if worker.is_connected else MeshNodeStatus.OFFLINE,
                tokens_per_second=worker.tokens_per_second or 0.0,
            )
            MeshRegistry.register_node(node)


def register_mac_agent_node(
    node_id: str,
    capabilities: List[str],
    supported_models: List[str] = None,
) -> None:
    node = MeshNode(
        node_id=node_id,
        node_type=NodeType.MAC_AGENT,
        runtime_type=WorkerRuntimeType.OLLAMA,
        capabilities=capabilities,
        supported_models=supported_models or ["llama3.2", "mistral"],
    )
    MeshRegistry.register_node(node)