import pytest
from unittest.mock import patch, MagicMock
import nomad_client
import subprocess

@patch("subprocess.run")
@patch("nomad_client.get_node_map")
def test_get_allocation_id(mock_node_map, mock_run):
    """Test getting allocation ID for a node."""
    mock_node_map.return_value = {"node_id1": "node1"}
    
    # Mock 'nomad job status navidrome-litefs' output
    mock_job_status = MagicMock()
    mock_job_status.stdout = """
Allocations
ID        Node ID   Task Group  Version  Desired  Status   Created    Modified
abc12345  node_id1  navidrome   1        run      running  1h ago     1h ago
"""
    
    # Mock 'nomad alloc status abc12345' output
    mock_alloc_status = MagicMock()
    mock_alloc_status.stdout = "ID = abc12345-full-id"
    
    mock_run.side_effect = [mock_job_status, mock_alloc_status]
    
    alloc_id = nomad_client.get_allocation_id("node1", "navidrome-litefs")
    assert alloc_id == "abc12345-full-id"

@patch("subprocess.run")
def test_get_logs(mock_run):
    """Test fetching logs for an allocation."""
    mock_stderr = "Error: database is locked\nSome other error"
    m = MagicMock()
    m.stdout = mock_stderr
    m.return_code = 0
    mock_run.return_value = m
    
    logs = nomad_client.get_allocation_logs("abc12345", tail=20)
    assert "database is locked" in logs
    # It should have tried with -task navidrome first
    mock_run.assert_any_call(
        ["nomad", "alloc", "logs", "-stderr", "-task", "navidrome", "-n", "20", "abc12345"],
        capture_output=True, text=True, check=True
    )

@patch("subprocess.run")
def test_restart_allocation(mock_run):
    """Test restarting an allocation."""
    m = MagicMock()
    m.return_code = 0
    mock_run.return_value = m
    
    success = nomad_client.restart_allocation("abc12345")
    assert success is True
    mock_run.assert_called_with(
        ["nomad", "alloc", "restart", "abc12345"],
        capture_output=True, text=True, check=True
    )

@patch("subprocess.run")
def test_exec_command(mock_run):
    """Test executing a command in an allocation."""
    m = MagicMock()
    m.stdout = "Command output"
    m.return_code = 0
    mock_run.return_value = m
    
    output = nomad_client.exec_command("abc12345", ["ls", "/data"])
    assert output == "Command output"
    mock_run.assert_called_with(
        ["nomad", "alloc", "exec", "-task", "navidrome", "abc12345", "ls", "/data"],
        capture_output=True, text=True, check=True
    )

@patch("subprocess.run")
def test_exec_command_failure(mock_run):
    """Test executing a command handles failure gracefully."""
    mock_run.side_effect = subprocess.CalledProcessError(1, "nomad", stderr="Nomad error")
    
    output = nomad_client.exec_command("abc12345", ["ls", "/data"])
    assert "Nomad Error" in output
    assert "Nomad error" not in output # The exception str might not contain stderr directly depending on python version
    
@patch("subprocess.run")
def test_get_node_map_failure(mock_run):
    """Test get_node_map handles failure."""
    mock_run.side_effect = FileNotFoundError("No such file")
    
    # It should not raise
    node_map = nomad_client.get_node_map()
    assert node_map == {}

@patch("subprocess.run")
def test_get_job_allocations(mock_run):
    """Test getting all allocations for a job with their IPs."""
    # Mock 'nomad job status navidrome-litefs'
    mock_job_status = MagicMock()
    mock_job_status.stdout = """
Allocations
ID        Node ID   Node Name  Status    Created
abc12345  node1     host1      running   1h ago
def67890  node2     host2      running   1h ago
"""
    
    # Mock 'nomad alloc status' for each alloc
    mock_alloc1 = MagicMock()
    mock_alloc1.stdout = """
ID = abc12345
Node Name = host1
Allocation Addresses:
Label    Dynamic  Address
*http    yes      1.1.1.1:4533 -> 4533
*litefs  yes      1.1.1.1:20202 -> 20202
"""
    mock_alloc2 = MagicMock()
    mock_alloc2.stdout = """
ID = def67890
Node Name = host2
Allocation Addresses:
Label    Dynamic  Address
*http    yes      2.2.2.2:4533 -> 4533
*litefs  yes      2.2.2.2:20202 -> 20202
"""
    
    mock_run.side_effect = [mock_job_status, mock_alloc1, mock_alloc2]
    
    # This should fail initially because nomad_client.get_job_allocations doesn't exist
    try:
        allocs = nomad_client.get_job_allocations("navidrome-litefs")
        assert len(allocs) == 2
        assert allocs[0]["ip"] == "1.1.1.1"
    except AttributeError:
        pytest.fail("nomad_client.get_job_allocations not implemented")