import json
import uuid
import pytest
import shutil
import tempfile
from pathlib import Path
from unittest.mock import patch, MagicMock

import brynq.compute_grid as cg
from brynq.compute_grid import (
    ComputeGrid,
    SwarmMapReduce,
    SwarmDAG,
    _topo_sort,
    register_map_function,
    register_reduce_function,
    reset_grid,
    list_jobs,
    job_history
)

@pytest.fixture(autouse=True)
def setup_test_env(monkeypatch):
    temp_dir = Path(tempfile.mkdtemp())
    monkeypatch.setattr(cg, "GRID_DIR", temp_dir)
    monkeypatch.setattr(cg, "NODES_FILE", temp_dir / "nodes.json")
    monkeypatch.setattr(cg, "JOBS_FILE", temp_dir / "jobs.json")
    monkeypatch.setattr(cg, "RESULTS_DIR", temp_dir / "results")
    monkeypatch.setattr(cg, "HISTORY_FILE", temp_dir / "history.jsonl")
    
    (temp_dir / "results").mkdir(parents=True, exist_ok=True)
    
    reset_grid()
    yield
    shutil.rmtree(temp_dir, ignore_errors=True)


def test_topo_sort_valid():
    dag = {"A": [], "B": ["A"], "C": ["A"], "D": ["B", "C"]}
    layers = _topo_sort(dag)
    assert layers[0] == ["A"]
    assert sorted(layers[1]) == ["B", "C"]
    assert layers[2] == ["D"]

def test_topo_sort_cycle():
    dag = {"A": ["B"], "B": ["A"]}
    with pytest.raises(ValueError, match="cycle"):
        _topo_sort(dag)

def test_register_functions():
    def dummy_map(x): return x
    def dummy_reduce(xs): return xs
    register_map_function("dummy_map", dummy_map)
    register_reduce_function("dummy_reduce", dummy_reduce)
    assert cg._map_functions["dummy_map"] == dummy_map
    assert cg._reduce_functions["dummy_reduce"] == dummy_reduce


class TestComputeGrid:
    def test_join_and_leave(self):
        node = ComputeGrid.join("test-agent", capabilities=["llm"])
        assert node["session_id"] == "test-agent"
        assert "llm" in node["capabilities"]
        assert node["status"] == "idle"
        
        nodes = ComputeGrid.list_nodes()
        assert len(nodes) == 1
        
        # Rejoin updates capabilities
        ComputeGrid.join("test-agent", capabilities=["llm", "search"], max_concurrent=10)
        nodes = ComputeGrid.list_nodes()
        assert len(nodes) == 1
        assert "search" in nodes[0]["capabilities"]
        assert nodes[0]["max_concurrent"] == 10
        
        left = ComputeGrid.leave("test-agent")
        assert left is True
        assert len(ComputeGrid.list_nodes()) == 0

    def test_heartbeat(self):
        ComputeGrid.join("agent-1")
        assert ComputeGrid.heartbeat("agent-1") is True
        assert ComputeGrid.heartbeat("nonexistent") is False

    def test_available_and_stats(self):
        ComputeGrid.join("agent-1", capabilities=["A"], max_concurrent=2)
        ComputeGrid.join("agent-2", capabilities=["B"], max_concurrent=2)
        
        available = ComputeGrid.available()
        assert len(available) == 2
        
        available_a = ComputeGrid.available(capability="A")
        assert len(available_a) == 1
        assert available_a[0]["session_id"] == "agent-1"
        
        stats = ComputeGrid.stats()
        assert stats["total_nodes"] == 2
        assert stats["total_capacity"] == 4
        
        # Assign task
        ComputeGrid._assign_task("agent-1")
        stats = ComputeGrid.stats()
        assert stats["active_tasks"] == 1
        
        ComputeGrid._assign_task("agent-1")
        assert ComputeGrid.available(capability="A") == []  # Reached max_concurrent=2
        
        ComputeGrid._release_task("agent-1")
        assert len(ComputeGrid.available(capability="A")) == 1


class TestSwarmMapReduce:
    def test_submit_invalid_functions(self):
        mr = SwarmMapReduce()
        with pytest.raises(ValueError, match="Map function.*not registered"):
            mr.submit("test", [1, 2], "missing", "missing")
            
        register_map_function("test_map", lambda x: x)
        with pytest.raises(ValueError, match="Reduce function.*not registered"):
            mr.submit("test", [1, 2], "test_map", "missing")

    def test_full_local_execution(self):
        # Setup functions
        def map_fn(x): return x * 2
        def reduce_fn(xs): return sum(xs)
        register_map_function("m_fn", map_fn)
        register_reduce_function("r_fn", reduce_fn)
        
        mr = SwarmMapReduce()
        job = mr.submit("test local", [1, 2, 3], "m_fn", "r_fn")
        job_id = job["job_id"]
        
        # Map Phase - should use local fallback since no nodes
        map_res = mr.map_phase(job_id)
        assert map_res["mapped"] == 3
        assert map_res["failed"] == 0
        assert "local" in map_res["nodes_used"]
        
        # Reduce Phase
        red_res = mr.reduce_phase(job_id)
        assert red_res["status"] == "completed"
        assert red_res["result"] == 12  # 1*2 + 2*2 + 3*2 = 12
        
        # List and history
        jobs = list_jobs()
        assert len(jobs) == 1
        assert jobs[0]["status"] == "completed"
        history = job_history()
        assert len(history) == 1

    @patch("brynq.compute_grid.service_mesh")
    def test_map_phase_with_pool_nodes(self, mock_sm):
        def map_fn(x): return x
        register_map_function("m", map_fn)
        register_reduce_function("r", lambda x: x)
        
        ComputeGrid.join("node-1", max_concurrent=10)
        
        # Mock service mesh call success
        mock_sm.call_service.return_value = {"success": True, "result": "mapped_val"}
        
        mr = SwarmMapReduce()
        job = mr.submit("test", [1, 2], "m", "r")
        job_id = job["job_id"]
        
        res = mr.map_phase(job_id)
        assert res["mapped"] == 2
        assert "node-1" in res["nodes_used"]
        
        status = mr.status(job_id)
        assert status["mapped"] == 2
        
    def test_cancel_job(self):
        register_map_function("m", lambda x: x)
        register_reduce_function("r", lambda x: x)
        
        mr = SwarmMapReduce()
        job = mr.submit("test", [1], "m", "r")
        res = mr.cancel(job["job_id"])
        assert res["status"] == "cancelled"
        
        # Already cancelled shouldn't be runnable
        with pytest.raises(ValueError, match="is in state 'cancelled'"):
            mr.map_phase(job["job_id"])
            
        # Re-submit for completion test
        job2 = mr.submit("test2", [1], "m", "r")
        mr.map_phase(job2["job_id"])
        mr.reduce_phase(job2["job_id"])
        with pytest.raises(ValueError):
            mr.cancel(job2["job_id"])


class TestSwarmDAG:
    def test_submit_and_execute_locally(self):
        def add_fn(data):
            c = data["chunk"]
            d = data["dependencies"]
            return c + sum(d.values())
            
        register_map_function("add", add_fn)
        
        dag = {
            "A": [],
            "B": ["A"],
            "C": ["A"],
            "D": ["B", "C"]
        }
        functions = {"A": "add", "B": "add", "C": "add", "D": "add"}
        payloads = {"A": 10, "B": 20, "C": 30, "D": 40}
        
        sd = SwarmDAG()
        job = sd.submit("dag test", dag, functions, payloads)
        job_id = job["job_id"]
        
        # Layer 0 (A)
        sd.execute_next_layer(job_id)
        j = cg._read_jobs()[job_id]
        assert j["results"]["A"] == 10
        assert j["current_layer"] == 1
        
        # Layer 1 (B, C)
        sd.execute_next_layer(job_id)
        j = cg._read_jobs()[job_id]
        assert j["results"]["B"] == 30  # 20 + 10
        assert j["results"]["C"] == 40  # 30 + 10
        assert j["current_layer"] == 2
        
        # Layer 2 (D)
        sd.execute_next_layer(job_id)
        j = cg._read_jobs()[job_id]
        assert j["results"]["D"] == 110 # 40 + 30 + 40
        assert j["status"] == "completed"
