"""
Tests for runpod.client module.

Tests GPU query parsing, pod deployment, error handling, and retry logic.
"""

import pytest
import responses
from unittest.mock import MagicMock, patch, Mock
import requests

from runpod.client import RunPodClient, RunPodAPIError
from runpod.config import RunPodConfig


class TestRunPodClient:
    """Tests for RunPodClient class."""

    def test_client_initialization(self, sample_config):
        """Test client initializes with config."""
        client = RunPodClient(sample_config)

        assert client.config == sample_config
        assert client.session is not None
        assert "Authorization" in client.session.headers
        assert client.session.headers["Authorization"] == f"Bearer {sample_config.api_key}"

    @responses.activate
    def test_query_graphql_success(self, sample_config, mock_gpu_data):
        """Test successful GraphQL query."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json=mock_gpu_data,
            status=200
        )

        client = RunPodClient(sample_config)
        result = client.query_graphql(RunPodClient.GPU_QUERY)

        assert result == mock_gpu_data
        assert len(responses.calls) == 1

    @responses.activate
    def test_query_graphql_with_variables(self, sample_config):
        """Test GraphQL query with variables."""
        expected_response = {"data": {"test": "value"}}
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json=expected_response,
            status=200
        )

        client = RunPodClient(sample_config)
        variables = {"gpuType": "RTX 4090"}
        result = client.query_graphql("query Test($gpuType: String)", variables)

        assert result == expected_response
        assert len(responses.calls) == 1

        # Verify variables were sent
        request_body = responses.calls[0].request.body
        assert b'"variables"' in request_body

    @responses.activate
    def test_query_graphql_error_response(self, sample_config, mock_api_responses):
        """Test GraphQL query with error response."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json=mock_api_responses["graphql_error"],
            status=200
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError, match="Invalid API key"):
            client.query_graphql(RunPodClient.GPU_QUERY)

    @responses.activate
    def test_query_graphql_network_error(self, sample_config):
        """Test GraphQL query with network error."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            body=requests.exceptions.ConnectionError("Network error")
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError, match="Failed to query"):
            client.query_graphql(RunPodClient.GPU_QUERY)

    @responses.activate
    def test_query_graphql_timeout(self, sample_config):
        """Test GraphQL query timeout."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            body=requests.exceptions.Timeout("Request timeout")
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError):
            client.query_graphql(RunPodClient.GPU_QUERY, timeout=1)

    @responses.activate
    def test_get_available_gpus_success(self, sample_config, mock_gpu_data):
        """Test getting available GPUs."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json=mock_gpu_data,
            status=200
        )

        client = RunPodClient(sample_config)
        gpus = client.get_available_gpus(min_vram_gb=16)

        # Should return 2 GPUs: RTX A4000 and Tesla V100 (secure cloud only)
        assert len(gpus) == 2

        # Should be sorted by price (V100 cheaper than A4000)
        assert gpus[0]["name"] == "Tesla V100"
        assert gpus[0]["price"] == 0.10
        assert gpus[1]["name"] == "RTX A4000"
        assert gpus[1]["price"] == 0.25

    @responses.activate
    def test_get_available_gpus_filters_vram(self, sample_config, mock_gpu_data):
        """Test GPU filtering by VRAM."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json=mock_gpu_data,
            status=200
        )

        client = RunPodClient(sample_config)
        gpus = client.get_available_gpus(min_vram_gb=24)

        # Only RTX 3090 has 24GB, but it's not in secure cloud
        assert len(gpus) == 0

    @responses.activate
    def test_get_available_gpus_filters_secure_cloud(self, sample_config, mock_gpu_data):
        """Test GPU filtering by secure cloud availability."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json=mock_gpu_data,
            status=200
        )

        client = RunPodClient(sample_config)
        gpus = client.get_available_gpus(min_vram_gb=16, secure_cloud_only=False)

        # Should include RTX 3090 (community cloud only)
        assert len(gpus) == 3  # A4000, V100, 3090

    @responses.activate
    def test_get_available_gpus_no_results(self, sample_config):
        """Test getting GPUs when none available."""
        responses.add(
            responses.POST,
            RunPodClient.GRAPHQL_ENDPOINT,
            json={"data": {"gpuTypes": []}},
            status=200
        )

        client = RunPodClient(sample_config)
        gpus = client.get_available_gpus()

        assert gpus == []

    @responses.activate
    def test_deploy_pod_success(self, sample_config, mock_pod_data):
        """Test successful pod deployment."""
        responses.add(
            responses.POST,
            RunPodClient.REST_API_URL,
            json=mock_pod_data,
            status=201
        )

        client = RunPodClient(sample_config)
        result = client.deploy_pod(
            gpu_id="NVIDIA RTX A4000",
            image="jgrusewski/foxhunt:latest",
            command="--epochs 50",
            container_disk_gb=50
        )

        assert result["id"] == "test-pod-123456"
        assert result["desiredStatus"] == "RUNNING"
        assert len(responses.calls) == 1

    @responses.activate
    def test_deploy_pod_with_custom_options(self, sample_config, mock_pod_data):
        """Test pod deployment with custom options."""
        responses.add(
            responses.POST,
            RunPodClient.REST_API_URL,
            json=mock_pod_data,
            status=200
        )

        client = RunPodClient(sample_config)
        result = client.deploy_pod(
            gpu_id="NVIDIA Tesla V100",
            image="custom/image:v2",
            command="python train.py --config custom",
            container_disk_gb=100,
            datacenters=["EUR-IS-1", "EUR-IS-2"],
            pod_name="custom-pod",
            ports=["8888/http", "6006/http", "22/tcp"]
        )

        assert result["id"] == "test-pod-123456"

        # Verify request payload
        request_body = responses.calls[0].request.body
        assert b'"name": "custom-pod"' in request_body
        assert b'"containerDiskInGb": 100' in request_body

    @responses.activate
    def test_deploy_pod_error_no_availability(self, sample_config, mock_api_responses):
        """Test pod deployment when GPU not available."""
        responses.add(
            responses.POST,
            RunPodClient.REST_API_URL,
            json=mock_api_responses["error_no_availability"],
            status=400
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError, match="No machines available"):
            client.deploy_pod(
                gpu_id="NVIDIA RTX A4000",
                image="test/image:latest"
            )

    @responses.activate
    def test_deploy_pod_error_invalid_response(self, sample_config):
        """Test pod deployment with invalid response format."""
        responses.add(
            responses.POST,
            RunPodClient.REST_API_URL,
            json={"invalid": "response"},  # Missing 'id' field
            status=200
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError, match="missing 'id' field"):
            client.deploy_pod(
                gpu_id="NVIDIA RTX A4000",
                image="test/image:latest"
            )

    @responses.activate
    def test_deploy_pod_network_error(self, sample_config):
        """Test pod deployment with network error."""
        responses.add(
            responses.POST,
            RunPodClient.REST_API_URL,
            body=requests.exceptions.ConnectionError("Network error")
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError, match="request failed"):
            client.deploy_pod(
                gpu_id="NVIDIA RTX A4000",
                image="test/image:latest"
            )

    @responses.activate
    def test_get_pod_status_success(self, sample_config, mock_api_responses):
        """Test getting pod status."""
        pod_id = "test-pod-123456"
        responses.add(
            responses.GET,
            f"{RunPodClient.REST_API_URL}/{pod_id}",
            json=mock_api_responses["pod_status_running"],
            status=200
        )

        client = RunPodClient(sample_config)
        status = client.get_pod_status(pod_id)

        assert status["id"] == pod_id
        assert status["desiredStatus"] == "RUNNING"

    @responses.activate
    def test_get_pod_status_error(self, sample_config):
        """Test getting pod status with error."""
        pod_id = "nonexistent-pod"
        responses.add(
            responses.GET,
            f"{RunPodClient.REST_API_URL}/{pod_id}",
            json={"error": "Pod not found"},
            status=404
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError):
            client.get_pod_status(pod_id)

    @responses.activate
    def test_stop_pod_success(self, sample_config):
        """Test stopping a pod."""
        pod_id = "test-pod-123456"
        responses.add(
            responses.DELETE,
            f"{RunPodClient.REST_API_URL}/{pod_id}",
            json={"success": True},
            status=200
        )

        client = RunPodClient(sample_config)
        result = client.stop_pod(pod_id)

        assert result is True

    @responses.activate
    def test_stop_pod_error(self, sample_config):
        """Test stopping pod with error."""
        pod_id = "test-pod-123456"
        responses.add(
            responses.DELETE,
            f"{RunPodClient.REST_API_URL}/{pod_id}",
            json={"error": "Failed to stop pod"},
            status=500
        )

        client = RunPodClient(sample_config)

        with pytest.raises(RunPodAPIError, match="Failed to stop"):
            client.stop_pod(pod_id)

    def test_deploy_pod_command_parsing(self, sample_config, mock_pod_data):
        """Test that command is properly split into arguments."""
        with patch('runpod.client.requests.Session') as mock_session_class:
            mock_session = MagicMock()
            mock_session_class.return_value = mock_session

            mock_response = MagicMock()
            mock_response.status_code = 201
            mock_response.json.return_value = mock_pod_data
            mock_session.post.return_value = mock_response

            client = RunPodClient(sample_config)

            client.deploy_pod(
                gpu_id="NVIDIA RTX A4000",
                image="test/image:latest",
                command='--parquet-file "/path/with spaces/file.parquet" --epochs 50'
            )

            # Verify the command was split correctly
            call_args = mock_session.post.call_args
            payload = call_args[1]['json']
            assert 'dockerStartCmd' in payload
            # Should be split into 4 arguments
            assert len(payload['dockerStartCmd']) == 4

    @pytest.mark.parametrize("status_code", [200, 201])
    def test_deploy_pod_accepts_both_success_codes(self, sample_config, mock_pod_data, status_code):
        """Test that both 200 and 201 are treated as success."""
        with responses.RequestsMock() as rsps:
            rsps.add(
                responses.POST,
                RunPodClient.REST_API_URL,
                json=mock_pod_data,
                status=status_code
            )

            client = RunPodClient(sample_config)
            result = client.deploy_pod(
                gpu_id="NVIDIA RTX A4000",
                image="test/image:latest"
            )

            assert result["id"] == "test-pod-123456"