"""
Phase 56: Compute Grid — Swarm MapReduce with DAG-based Distributed Execution.

A decentralized compute grid that allows agents to pool resources across the
federated network. Combines a ComputeGrid node manager with a SwarmMapReduce
coordinator that breaks parent tasks into DAGs of micro-tasks, fans them out
to idle compute nodes, and reduces results into a single answer.

Integration:
  - service_mesh: registers "compute_grid" service, discovers compute nodes
  - task_queue: enqueues micro-tasks for agent pickup
  - swarm: leverages swarm formation for worker recruitment

Storage:
  state/broker/compute_grid/
    nodes.json               — registered compute pool nodes
    jobs.json                — active job registry
    results/{job_id}.jsonl   — per-job shard results
    history.jsonl            — completed / cancelled job log
"""

import json
import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Set

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync
from brynq import service_mesh
from brynq import marketplace_payments

logger = logging.getLogger(__name__)

# ── Paths ──────────────────────────────────────────────────────────────────

GRID_DIR = BROKER_DIR / "compute_grid"
NODES_FILE = GRID_DIR / "nodes.json"
JOBS_FILE = GRID_DIR / "jobs.json"
RESULTS_DIR = GRID_DIR / "results"
HISTORY_FILE = GRID_DIR / "history.jsonl"

for _d in (GRID_DIR, RESULTS_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# ── In-memory registries ─────────────────────────────────────────────────

_map_functions: Dict[str, Callable] = {}
_reduce_functions: Dict[str, Callable] = {}


def reset_grid() -> None:
    """Reset all in-memory state. Used by tests."""
    global _map_functions, _reduce_functions
    _map_functions = {}
    _reduce_functions = {}


# ── Internal helpers ─────────────────────────────────────────────────────

def _read_nodes() -> List[dict]:
    """Load compute pool nodes from disk."""
    if not NODES_FILE.exists():
        return []
    try:
        return json.loads(NODES_FILE.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return []


def _save_nodes(nodes: List[dict]) -> None:
    _write_json_sync(NODES_FILE, nodes)


def _read_jobs() -> Dict[str, dict]:
    """Load active jobs from disk."""
    if not JOBS_FILE.exists():
        return {}
    try:
        return json.loads(JOBS_FILE.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return {}


def _save_jobs(jobs: Dict[str, dict]) -> None:
    """Persist jobs to disk."""
    _write_json_sync(JOBS_FILE, jobs)


def _results_file(job_id: str) -> Path:
    return RESULTS_DIR / f"{job_id}.jsonl"


def _read_results(job_id: str) -> List[dict]:
    return _read_jsonl_sync(_results_file(job_id))


def _append_result(job_id: str, result: dict) -> None:
    _append_jsonl_sync(_results_file(job_id), result)


# ── DAG helpers ──────────────────────────────────────────────────────────

def _topo_sort(dag: Dict[str, List[str]]) -> List[List[str]]:
    """Topological sort returning execution layers (parallelizable groups).

    Each layer contains nodes whose dependencies are fully satisfied by
    prior layers. Nodes within a layer can execute in parallel.

    Args:
        dag: Dict mapping node_id -> list of dependency node_ids.

    Returns:
        List of layers, each layer a list of node_ids.

    Raises:
        ValueError: If the DAG contains a cycle.
    """
    in_degree: Dict[str, int] = {node: len(deps) for node, deps in dag.items()}
    reverse: Dict[str, List[str]] = {node: [] for node in dag}
    for node, deps in dag.items():
        for dep in deps:
            if dep not in reverse:
                reverse[dep] = []
            reverse[dep].append(node)

    layers: List[List[str]] = []
    ready = [n for n, d in in_degree.items() if d == 0]

    visited = 0
    while ready:
        layers.append(sorted(ready))
        next_ready = []
        for node in ready:
            visited += 1
            for dependent in reverse.get(node, []):
                in_degree[dependent] -= 1
                if in_degree[dependent] == 0:
                    next_ready.append(dependent)
        ready = next_ready

    if visited < len(dag):
        raise ValueError("DAG contains a cycle")

    return layers


# ── Function Registry ────────────────────────────────────────────────────

def register_map_function(name: str, fn: Callable) -> None:
    """Register a named map function for use in jobs.

    The function receives a single chunk (any type) and returns a result dict.
    """
    _map_functions[name] = fn
    logger.info("Registered map function '%s'", name)


def register_reduce_function(name: str, fn: Callable) -> None:
    """Register a named reduce function for use in jobs.

    The function receives a list of shard result dicts and returns a
    single combined result dict.
    """
    _reduce_functions[name] = fn
    logger.info("Registered reduce function '%s'", name)


# ── ComputeGrid ──────────────────────────────────────────────────────────

class ComputeGrid:
    """Manages the pool of idle compute nodes across the federated network.

    Agents toggle into "compute mode" by joining the grid. The grid tracks
    their capacity, capabilities, and heartbeat status. SwarmMapReduce uses
    this pool to fan out micro-tasks.
    """

    @staticmethod
    def join(
        session_id: str,
        capabilities: Optional[List[str]] = None,
        max_concurrent: int = 4,
    ) -> dict:
        """An agent joins the compute pool.

        Args:
            session_id: The agent's broker session ID.
            capabilities: Capability tags (e.g. ["llm", "code_exec", "search"]).
            max_concurrent: Max micro-tasks this node handles at once.

        Returns:
            The node record.
        """
        nodes = _read_nodes()
        now = datetime.now(timezone.utc).isoformat()

        # Update if already registered
        for node in nodes:
            if node["session_id"] == session_id:
                node["capabilities"] = capabilities or node.get("capabilities", [])
                node["max_concurrent"] = max_concurrent
                node["status"] = "idle"
                node["last_heartbeat"] = now
                _save_nodes(nodes)
                logger.info("Compute node %s rejoined pool", session_id)
                return node

        node = {
            "session_id": session_id,
            "capabilities": capabilities or [],
            "max_concurrent": max_concurrent,
            "active_tasks": 0,
            "completed_tasks": 0,
            "failed_tasks": 0,
            "status": "idle",
            "joined_at": now,
            "last_heartbeat": now,
        }
        nodes.append(node)
        _save_nodes(nodes)
        logger.info("Compute node %s joined pool", session_id)
        return node

    @staticmethod
    def leave(session_id: str) -> bool:
        """Remove a node from the compute pool.

        Returns:
            True if found and removed.
        """
        nodes = _read_nodes()
        before = len(nodes)
        nodes = [n for n in nodes if n["session_id"] != session_id]
        if len(nodes) == before:
            return False
        _save_nodes(nodes)
        logger.info("Compute node %s left pool", session_id)
        return True

    @staticmethod
    def heartbeat(session_id: str) -> bool:
        """Update a node's heartbeat. Returns True if found."""
        nodes = _read_nodes()
        for node in nodes:
            if node["session_id"] == session_id:
                node["last_heartbeat"] = datetime.now(timezone.utc).isoformat()
                _save_nodes(nodes)
                return True
        return False

    @staticmethod
    def list_nodes(status: Optional[str] = None) -> List[dict]:
        """List all nodes in the compute pool.

        Args:
            status: Optional filter: "idle", "busy", "offline".
        """
        nodes = _read_nodes()
        if status is not None:
            nodes = [n for n in nodes if n.get("status") == status]
        return nodes

    @staticmethod
    def available(capability: Optional[str] = None) -> List[dict]:
        """Find nodes with spare capacity, optionally matching a capability."""
        nodes = _read_nodes()
        available = []
        for node in nodes:
            if node.get("status") == "offline":
                continue
            if node.get("active_tasks", 0) >= node.get("max_concurrent", 4):
                continue
            if capability and capability not in node.get("capabilities", []):
                continue
            available.append(node)
        return available

    @staticmethod
    def stats() -> dict:
        """Return compute grid pool statistics."""
        nodes = _read_nodes()
        total = len(nodes)
        idle = sum(1 for n in nodes if n.get("status") == "idle")
        busy = sum(1 for n in nodes if n.get("status") == "busy")
        capacity = sum(n.get("max_concurrent", 4) for n in nodes)
        active = sum(n.get("active_tasks", 0) for n in nodes)

        history = _read_jsonl_sync(HISTORY_FILE)
        return {
            "total_nodes": total,
            "idle_nodes": idle,
            "busy_nodes": busy,
            "total_capacity": capacity,
            "active_tasks": active,
            "available_capacity": capacity - active,
            "jobs_completed": len(history),
        }

    @staticmethod
    def _assign_task(session_id: str) -> bool:
        """Increment a node's active_tasks count. Returns False if at capacity."""
        nodes = _read_nodes()
        for node in nodes:
            if node["session_id"] == session_id:
                if node.get("active_tasks", 0) >= node.get("max_concurrent", 4):
                    return False
                node["active_tasks"] = node.get("active_tasks", 0) + 1
                if node["active_tasks"] >= node.get("max_concurrent", 4):
                    node["status"] = "busy"
                _save_nodes(nodes)
                return True
        return False

    @staticmethod
    def _release_task(session_id: str, failed: bool = False) -> None:
        """Decrement a node's active_tasks after completion."""
        nodes = _read_nodes()
        for node in nodes:
            if node["session_id"] == session_id:
                node["active_tasks"] = max(0, node.get("active_tasks", 0) - 1)
                if failed:
                    node["failed_tasks"] = node.get("failed_tasks", 0) + 1
                else:
                    node["completed_tasks"] = node.get("completed_tasks", 0) + 1
                if node["active_tasks"] < node.get("max_concurrent", 4):
                    node["status"] = "idle"
                _save_nodes(nodes)
                return


# ── SwarmMapReduce ───────────────────────────────────────────────────────

class SwarmMapReduce:
    """Coordinator for distributed map-reduce jobs across the compute mesh.

    Workflow:
        1. submit() — register a job with chunks
        2. map_phase() — fan out chunks to compute nodes
        3. reduce_phase() — collect results and synthesize
    """

    def submit(
        self,
        task_description: str,
        chunks: List[Any],
        map_fn_name: str,
        reduce_fn_name: str,
        dag: Optional[Dict[str, List[str]]] = None,
        capability_required: Optional[str] = None,
    ) -> dict:
        """Create a new MapReduce job.

        Args:
            task_description: Human-readable description of the job.
            chunks: List of data chunks to distribute.
            map_fn_name: Name of a registered map function.
            reduce_fn_name: Name of a registered reduce function.
            dag: Optional DAG defining dependencies between chunk indices.
                 Keys are chunk index strings ("0", "1", ...), values are
                 lists of indices that must complete first. If None, all
                 chunks run in a single parallel layer.
            capability_required: Optional capability that compute nodes must
                                 have to accept work from this job.

        Returns:
            The job record.
        """
        if map_fn_name not in _map_functions:
            raise ValueError(f"Map function '{map_fn_name}' not registered")
        if reduce_fn_name not in _reduce_functions:
            raise ValueError(f"Reduce function '{reduce_fn_name}' not registered")
        if not chunks:
            raise ValueError("chunks list must not be empty")

        job_id = uuid.uuid4().hex[:12]
        now = datetime.now(timezone.utc).isoformat()

        # Build DAG layers for execution ordering
        if dag is not None:
            layers = _topo_sort(dag)
        else:
            # All chunks in a single parallel layer
            layers = [[str(i) for i in range(len(chunks))]]

        job = {
            "job_id": job_id,
            "description": task_description,
            "map_fn": map_fn_name,
            "reduce_fn": reduce_fn_name,
            "chunk_count": len(chunks),
            "chunks": chunks,
            "dag": dag,
            "layers": layers,
            "current_layer": 0,
            "capability_required": capability_required,
            "status": "submitted",
            "created_at": now,
            "updated_at": now,
            "mapped": 0,
            "failed": 0,
            "reduced": False,
            "result": None,
            "nodes_used": [],
        }

        jobs = _read_jobs()
        jobs[job_id] = job
        _save_jobs(jobs)

        logger.info("Submitted MapReduce job '%s' with %d chunks in %d layers", job_id, len(chunks), len(layers))
        return job

    def map_phase(self, job_id: str) -> dict:
        """Execute the map phase: distribute chunks to compute nodes.

        Executes DAG layers sequentially. Within each layer, chunks are
        distributed in parallel across available compute pool nodes (via
        ComputeGrid). Falls back to service_mesh discovery, then to local
        execution if no pool or mesh nodes are available.

        Returns:
            Summary dict with mapped/failed counts per layer.
        """
        jobs = _read_jobs()
        job = jobs.get(job_id)
        if job is None:
            raise ValueError(f"Job '{job_id}' not found")
        if job["status"] not in ("submitted", "mapping"):
            raise ValueError(f"Job '{job_id}' is in state '{job['status']}', expected 'submitted' or 'mapping'")

        map_fn = _map_functions.get(job["map_fn"])
        if map_fn is None:
            raise ValueError(f"Map function '{job['map_fn']}' no longer registered")

        job["status"] = "mapping"
        job["updated_at"] = datetime.now(timezone.utc).isoformat()
        _save_jobs(jobs)

        layers = job.get("layers", [[str(i) for i in range(job["chunk_count"])]])
        capability = job.get("capability_required")
        chunks = job["chunks"]
        all_nodes_used: list = []

        mapped = 0
        failed = 0
        layer_results: List[dict] = []

        for layer_idx, layer in enumerate(layers):
            layer_mapped = 0
            layer_failed = 0

            # Resolve available nodes: prefer ComputeGrid pool, then service_mesh
            pool_nodes = ComputeGrid.available(capability=capability)
            if pool_nodes:
                node_targets = [n["session_id"] for n in pool_nodes]
                node_source = "pool"
            else:
                mesh_nodes = service_mesh.discover(capability="compute")
                node_targets = [n["name"] for n in mesh_nodes]
                node_source = "mesh" if node_targets else "local"

            for idx_in_layer, chunk_key in enumerate(layer):
                chunk_index = int(chunk_key)
                chunk = chunks[chunk_index]

                shard = {
                    "shard_index": chunk_index,
                    "layer": layer_idx,
                    "chunk": chunk,
                    "status": "pending",
                    "node": None,
                    "node_source": node_source,
                    "started_at": datetime.now(timezone.utc).isoformat(),
                }

                try:
                    if node_source == "pool":
                        target_sid = node_targets[idx_in_layer % len(node_targets)]
                        ComputeGrid._assign_task(target_sid)
                        try:
                            # Actually dispatch via service_mesh using target_sid as the node name
                            result = service_mesh.call_service(
                                target_sid, "execute_map",
                                {"map_fn": job["map_fn"], "chunk": chunk, "job_id": job_id, "shard_index": chunk_index},
                            )
                            if result.get("success"):
                                shard["status"] = "completed"
                                shard["result"] = result.get("result")
                                layer_mapped += 1
                                # Reward the node for compute time
                                marketplace_payments._record_tx(
                                    agent_id=target_sid,
                                    tx_type="task_earned",
                                    amount=0.01,
                                    description=f"Leased compute for job {job_id}",
                                    reference_id=job_id
                                )
                            else:
                                shard["status"] = "failed"
                                shard["error"] = result.get("error", "unknown")
                                layer_failed += 1
                        except Exception as e:
                            shard["status"] = "failed"
                            shard["error"] = str(e)
                            layer_failed += 1
                        finally:
                            ComputeGrid._release_task(target_sid, failed=shard["status"] == "failed")
                        shard["node"] = target_sid

                    elif node_source == "mesh":
                        target_node = node_targets[idx_in_layer % len(node_targets)]
                        result = service_mesh.call_service(
                            target_node, "execute_map",
                            {"map_fn": job["map_fn"], "chunk": chunk, "job_id": job_id, "shard_index": chunk_index},
                        )
                        shard["node"] = target_node
                        if result.get("success"):
                            shard["status"] = "completed"
                            shard["result"] = result.get("result")
                            layer_mapped += 1
                            # Reward the node for compute time
                            marketplace_payments._record_tx(
                                agent_id=target_node,
                                tx_type="task_earned",
                                amount=0.01,
                                description=f"Leased compute for job {job_id}",
                                reference_id=job_id
                            )
                        else:
                            shard["status"] = "failed"
                            shard["error"] = result.get("error", "unknown")
                            layer_failed += 1

                    else:
                        # Local fallback
                        shard["node"] = "local"
                        shard_result = map_fn(chunk)
                        shard["status"] = "completed"
                        shard["result"] = shard_result
                        layer_mapped += 1

                except Exception as e:
                    shard["status"] = "failed"
                    shard["error"] = str(e)
                    layer_failed += 1

                shard["completed_at"] = datetime.now(timezone.utc).isoformat()
                _append_result(job_id, shard)

            mapped += layer_mapped
            failed += layer_failed
            all_nodes_used.extend(node_targets if node_targets else ["local"])
            layer_results.append({"layer": layer_idx, "mapped": layer_mapped, "failed": layer_failed})

        # Update job
        jobs = _read_jobs()
        job = jobs[job_id]
        job["mapped"] = mapped
        job["failed"] = failed
        job["nodes_used"] = list(set(all_nodes_used))
        job["current_layer"] = len(layers)
        job["status"] = "mapped"
        job["updated_at"] = datetime.now(timezone.utc).isoformat()
        _save_jobs(jobs)

        logger.info("Map phase complete for job '%s': %d mapped, %d failed across %d layers", job_id, mapped, failed, len(layers))
        return {"job_id": job_id, "mapped": mapped, "failed": failed, "layers": layer_results, "nodes_used": job["nodes_used"]}

    def reduce_phase(self, job_id: str) -> dict:
        """Execute the reduce phase: collect shard results and synthesize.

        Returns:
            The final reduced result.
        """
        jobs = _read_jobs()
        job = jobs.get(job_id)
        if job is None:
            raise ValueError(f"Job '{job_id}' not found")
        if job["status"] != "mapped":
            raise ValueError(f"Job '{job_id}' is in state '{job['status']}', expected 'mapped'")

        reduce_fn = _reduce_functions.get(job["reduce_fn"])
        if reduce_fn is None:
            raise ValueError(f"Reduce function '{job['reduce_fn']}' no longer registered")

        job["status"] = "reducing"
        job["updated_at"] = datetime.now(timezone.utc).isoformat()
        _save_jobs(jobs)

        # Collect successful shard results
        shards = _read_results(job_id)
        shard_results = [s["result"] for s in shards if s.get("status") == "completed" and "result" in s]

        try:
            final_result = reduce_fn(shard_results)
            job["result"] = final_result
            job["reduced"] = True
            job["status"] = "completed"
        except Exception as e:
            job["status"] = "reduce_failed"
            job["result"] = {"error": str(e)}
            logger.error("Reduce failed for job '%s': %s", job_id, e)

        job["updated_at"] = datetime.now(timezone.utc).isoformat()

        # Persist and archive
        jobs = _read_jobs()
        jobs[job_id] = job
        _save_jobs(jobs)
        _append_jsonl_sync(HISTORY_FILE, job)

        logger.info("Job '%s' finished with status '%s'", job_id, job["status"])
        return {"job_id": job_id, "status": job["status"], "result": job["result"]}

    def status(self, job_id: str) -> dict:
        """Get current status of a job.

        Returns:
            Dict with job metadata and progress counters.
        """
        jobs = _read_jobs()
        job = jobs.get(job_id)
        if job is None:
            raise ValueError(f"Job '{job_id}' not found")

        return {
            "job_id": job_id,
            "description": job["description"],
            "status": job["status"],
            "chunk_count": job["chunk_count"],
            "mapped": job["mapped"],
            "failed": job["failed"],
            "reduced": job["reduced"],
            "nodes_used": job["nodes_used"],
            "created_at": job["created_at"],
            "updated_at": job["updated_at"],
        }

    def cancel(self, job_id: str) -> dict:
        """Cancel a job.

        Returns:
            Dict confirming cancellation.
        """
        jobs = _read_jobs()
        job = jobs.get(job_id)
        if job is None:
            raise ValueError(f"Job '{job_id}' not found")
        if job["status"] == "completed":
            raise ValueError(f"Job '{job_id}' is already completed")

        job["status"] = "cancelled"
        job["updated_at"] = datetime.now(timezone.utc).isoformat()
        _save_jobs(jobs)
        _append_jsonl_sync(HISTORY_FILE, job)

        logger.info("Cancelled job '%s'", job_id)
        return {"job_id": job_id, "status": "cancelled"}


# ── Module-level convenience ─────────────────────────────────────────────

def list_jobs(status_filter: Optional[str] = None) -> List[dict]:
    """List all active jobs, optionally filtered by status."""
    jobs = _read_jobs()
    results = []
    for job in jobs.values():
        if status_filter and job.get("status") != status_filter:
            continue
        results.append({
            "job_id": job["job_id"],
            "description": job["description"],
            "status": job["status"],
            "chunk_count": job["chunk_count"],
            "mapped": job["mapped"],
            "failed": job["failed"],
            "created_at": job["created_at"],
        })
    return sorted(results, key=lambda j: j["created_at"], reverse=True)


def job_history(limit: int = 50) -> List[dict]:
    """Return recent completed/cancelled jobs from history."""
    entries = _read_jsonl_sync(HISTORY_FILE)
    return entries[-limit:]


# ── Compute Leasing (The Economy) ────────────────────────────────────────

COMPUTE_PRICE_PER_SHARD = 1.0  # credits per micro-task/shard executed


def lease_compute(job_id: str, payer_id: str) -> dict:
    """Create an escrow payment for a compute job.

    Calculates cost based on chunk_count * COMPUTE_PRICE_PER_SHARD and
    locks credits from the payer via marketplace_payments escrow.

    Args:
        job_id: The MapReduce job to pay for.
        payer_id: The agent session ID submitting the job.

    Returns:
        Dict with escrow details and total cost, or error dict.
    """
    from brynq import marketplace_payments

    jobs = _read_jobs()
    job = jobs.get(job_id)
    if job is None:
        return {"error": f"Job '{job_id}' not found"}

    total_cost = round(job["chunk_count"] * COMPUTE_PRICE_PER_SHARD, 2)
    if total_cost <= 0:
        return {"error": "Job has no chunks to pay for"}

    escrow = marketplace_payments.create_escrow(
        task_id=f"compute_{job_id}",
        payer_id=payer_id,
        amount=total_cost,
        description=f"Compute lease for job {job_id} ({job['chunk_count']} shards @ {COMPUTE_PRICE_PER_SHARD} credits/shard)",
    )

    if "error" in escrow:
        return escrow

    # Tag the job with payment info
    job["payment"] = {
        "escrow_task_id": f"compute_{job_id}",
        "payer_id": payer_id,
        "total_cost": total_cost,
        "price_per_shard": COMPUTE_PRICE_PER_SHARD,
        "status": "escrowed",
    }
    _save_jobs(jobs)

    logger.info("Compute lease created for job %s: %s credits escrowed", job_id, total_cost)
    return {
        "job_id": job_id,
        "total_cost": total_cost,
        "payer_id": payer_id,
        "escrow": escrow,
    }


def settle_compute(job_id: str) -> dict:
    """Release escrowed payment to compute nodes that processed shards.

    Distributes credits proportionally to nodes based on completed shards.
    Failed shards are not paid. Platform fee is taken per marketplace_payments
    standard rate.

    Args:
        job_id: The completed MapReduce job to settle.

    Returns:
        Dict with settlement details per node, or error dict.
    """
    from brynq import marketplace_payments

    jobs = _read_jobs()
    job = jobs.get(job_id)
    if job is None:
        return {"error": f"Job '{job_id}' not found"}

    payment = job.get("payment")
    if not payment:
        return {"error": f"Job '{job_id}' has no payment/escrow"}
    if payment.get("status") == "settled":
        return {"error": f"Job '{job_id}' already settled"}

    if job["status"] not in ("completed", "mapped", "reducing"):
        return {"error": f"Job '{job_id}' not yet complete (status: {job['status']})"}

    # Count completed shards per node
    shards = _read_results(job_id)
    node_shard_counts: Dict[str, int] = {}
    for shard in shards:
        if shard.get("status") == "completed":
            node = shard.get("node", "local")
            node_shard_counts[node] = node_shard_counts.get(node, 0) + 1

    if not node_shard_counts:
        # All shards failed — refund
        marketplace_payments.refund_escrow(f"compute_{job_id}")
        payment["status"] = "refunded"
        _save_jobs(jobs)
        return {"job_id": job_id, "status": "refunded", "reason": "all shards failed"}

    # Release escrow — the platform takes its fee, rest goes to workers
    total_completed = sum(node_shard_counts.values())
    price_per_shard = payment.get("price_per_shard", COMPUTE_PRICE_PER_SHARD)
    payable_amount = round(total_completed * price_per_shard, 2)

    # Pay each node proportionally
    settlements = []
    for node_id, count in node_shard_counts.items():
        if node_id == "local":
            continue  # Local execution doesn't earn credits
        node_pay = round(count * price_per_shard, 2)
        marketplace_payments._record_tx(
            agent_id=node_id,
            tx_type="task_earned",
            amount=node_pay,
            description=f"Compute lease payment: {count} shards @ {price_per_shard}/shard for job {job_id}",
            reference_id=f"compute_{job_id}",
        )
        settlements.append({"node_id": node_id, "shards": count, "paid": node_pay})

    # Record platform fee
    platform_fee = round(payable_amount * marketplace_payments.PLATFORM_FEE_RATE, 2)
    if platform_fee > 0:
        marketplace_payments._record_tx(
            agent_id="brynq_platform",
            tx_type="platform_fee",
            amount=platform_fee,
            description=f"Compute grid fee for job {job_id}",
            reference_id=f"compute_{job_id}",
        )

    # Mark escrow as released
    escrow_path = marketplace_payments.ESCROW_DIR / f"compute_{job_id}.json"
    try:
        escrow = json.loads(escrow_path.read_text("utf-8"))
        escrow["status"] = "released"
        escrow["released_at"] = datetime.now(timezone.utc).isoformat()
        marketplace_payments._write_json_sync(escrow_path, escrow)
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        pass

    payment["status"] = "settled"
    payment["settlements"] = settlements
    payment["platform_fee"] = platform_fee
    _save_jobs(jobs)

    logger.info("Compute lease settled for job %s: %d nodes paid, %s platform fee", job_id, len(settlements), platform_fee)
    return {
        "job_id": job_id,
        "total_paid": payable_amount,
        "platform_fee": platform_fee,
        "settlements": settlements,
    }


# ── SwarmDAG ─────────────────────────────────────────────────────────────

class SwarmDAG:
    """Coordinator for executing arbitrary DAGs of tasks on the compute grid.
    
    Allows breaking down a massive task into hundreds of dependent micro-tasks,
    resolving execution order via topological sort, and dispatching layers in
    parallel to the Swarm compute mesh.
    """

    def submit(
        self,
        description: str,
        dag: Dict[str, List[str]],
        functions: Dict[str, str],
        payloads: Dict[str, Any],
    ) -> dict:
        """Create a new DAG distributed job.

        Args:
            description: Human-readable job description.
            dag: Dict mapping node_id -> list of dependency node_ids.
            functions: Dict mapping node_id -> registered map/reduce function name.
            payloads: Dict mapping node_id -> data payload/chunk.

        Returns:
            The job record.
        """
        job_id = uuid.uuid4().hex[:12]
        now = datetime.now(timezone.utc).isoformat()

        try:
            layers = _topo_sort(dag)
        except ValueError as e:
            raise ValueError(f"Invalid DAG: {e}")

        job = {
            "job_id": job_id,
            "type": "dag",
            "description": description,
            "dag": dag,
            "layers": layers,
            "functions": functions,
            "payloads": payloads,
            "status": "submitted",
            "current_layer": 0,
            "results": {},
            "nodes_used": [],
            "created_at": now,
            "updated_at": now,
        }

        jobs = _read_jobs()
        jobs[job_id] = job
        _save_jobs(jobs)

        logger.info("Submitted DAG job '%s' with %d layers", job_id, len(layers))
        return job

    def execute_next_layer(self, job_id: str) -> dict:
        """Execute the next ready layer of the DAG across the compute mesh.
        
        Gathers results from dependencies, injects them into the current layer's
        payloads, and fans out execution.
        """
        jobs = _read_jobs()
        job = jobs.get(job_id)
        if not job or job.get("type") != "dag":
            raise ValueError(f"DAG Job '{job_id}' not found")
        
        if job["status"] in ("completed", "failed", "cancelled"):
            return job
            
        layer_idx = job["current_layer"]
        layers = job["layers"]
        if layer_idx >= len(layers):
            job["status"] = "completed"
            job["updated_at"] = datetime.now(timezone.utc).isoformat()
            _save_jobs(jobs)
            _append_jsonl_sync(HISTORY_FILE, job)
            return job
            
        current_nodes = layers[layer_idx]
        job["status"] = f"executing_layer_{layer_idx}"
        _save_jobs(jobs)
        
        # Discover compute nodes
        nodes = service_mesh.discover(capability="compute")
        node_names = [n["name"] for n in nodes]
        
        failed = 0
        layer_results = {}
        used_nodes = set(job.get("nodes_used", []))
        
        for i, node_id in enumerate(current_nodes):
            fn_name = job["functions"].get(node_id)
            payload = job["payloads"].get(node_id)
            
            # Inject dependencies' results into the payload
            deps = job["dag"].get(node_id, [])
            dep_results = {d: job["results"].get(d) for d in deps}
            
            call_payload = {
                "chunk": payload,
                "dependencies": dep_results
            }

            try:
                if node_names:
                    target = node_names[i % len(node_names)]
                    used_nodes.add(target)
                    result = service_mesh.call_service(
                        target, "execute_map",
                        {"map_fn": fn_name, "chunk": call_payload, "job_id": job_id, "shard_index": node_id}
                    )
                    if result.get("success"):
                        layer_results[node_id] = result.get("result")
                        # Reward the node for compute time
                        marketplace_payments._record_tx(
                            agent_id=target,
                            tx_type="task_earned",
                            amount=0.01,
                            description=f"Leased compute for DAG job {job_id} node {node_id}",
                            reference_id=job_id
                        )
                    else:
                        failed += 1
                        logger.error("DAG node %s failed: %s", node_id, result.get("error"))
                else:
                    used_nodes.add("local")
                    fn = _map_functions.get(fn_name) or _reduce_functions.get(fn_name)
                    if fn:
                        res = fn(call_payload)
                        layer_results[node_id] = res
                    else:
                        failed += 1
                        logger.error("DAG node %s failed: function %s not registered", node_id, fn_name)
            except Exception as e:
                logger.error("DAG node %s failed locally: %s", node_id, e)
                failed += 1
                
        job["nodes_used"] = list(used_nodes)
        
        if failed > 0:
            job["status"] = "failed"
        else:
            job["results"].update(layer_results)
            job["current_layer"] += 1
            if job["current_layer"] >= len(layers):
                job["status"] = "completed"
                _append_jsonl_sync(HISTORY_FILE, job)
                
        job["updated_at"] = datetime.now(timezone.utc).isoformat()
        _save_jobs(jobs)
        return job
