import json
from brynq import compute_grid

def mcp_broker_cg_reset(args: str) -> str:
    """Reset the compute grid state."""
    compute_grid.reset_grid()
    return "Compute grid reset successfully."

def mcp_broker_cg_submit(args: str) -> str:
    """Submit a map-reduce job to the compute grid. Args: JSON containing job_type, inputs (list), and optional dag (dict)."""
    try:
        data = json.loads(args)
        job_id = compute_grid.submit_job(data.get('job_type', 'default'), data.get('inputs', []), data.get('dag'))
        return json.dumps({"status": "success", "job_id": job_id})
    except Exception as e:
        return json.dumps({"error": str(e)})

def mcp_broker_cg_status(args: str) -> str:
    """Get status of a compute grid job. Args: job_id"""
    job_id = args.strip()
    if not job_id:
        return "Error: Missing job_id"
    status = compute_grid.get_job_status(job_id)
    return json.dumps(status)

TOOLS = [
    {
        "name": "broker_cg_reset",
        "description": "Reset the compute grid state.",
        "inputSchema": {"type": "object", "properties": {"args": {"type": "string"}}},
        "handler": mcp_broker_cg_reset
    },
    {
        "name": "broker_cg_submit",
        "description": "Submit a map-reduce job. Args: JSON containing job_type, inputs, dag.",
        "inputSchema": {"type": "object", "properties": {"args": {"type": "string"}}},
        "handler": mcp_broker_cg_submit
    },
    {
        "name": "broker_cg_status",
        "description": "Get status of a compute grid job. Args: job_id",
        "inputSchema": {"type": "object", "properties": {"args": {"type": "string"}}},
        "handler": mcp_broker_cg_status
    }
]
