import asyncio
import json
from dataclasses import asdict
from unittest.mock import AsyncMock, MagicMock

import pytest
from google.protobuf.json_format import Parse

from ray.core.generated.gcs_pb2 import JobsAPIInfo
from ray.dashboard.modules.job.common import (
    JobErrorType,
    JobInfo,
    JobInfoStorageClient,
    JobStatus,
    JobSubmitRequest,
    http_uri_components_to_uri,
    uri_to_http_components,
    validate_request_type,
)


class TestJobSubmitRequestValidation:
    def test_validate_entrypoint(self):
        r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
        assert r.entrypoint == "abc"

        with pytest.raises(TypeError, match="required positional argument"):
            validate_request_type({}, JobSubmitRequest)

        with pytest.raises(TypeError, match="must be a string"):
            validate_request_type({"entrypoint": 123}, JobSubmitRequest)

    def test_validate_submission_id(self):
        r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
        assert r.entrypoint == "abc"
        assert r.submission_id is None

        r = validate_request_type(
            {"entrypoint": "abc", "submission_id": "123"}, JobSubmitRequest
        )
        assert r.entrypoint == "abc"
        assert r.submission_id == "123"

        with pytest.raises(TypeError, match="must be a string"):
            validate_request_type(
                {"entrypoint": 123, "submission_id": 1}, JobSubmitRequest
            )

    def test_validate_runtime_env(self):
        r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
        assert r.entrypoint == "abc"
        assert r.runtime_env is None

        r = validate_request_type(
            {"entrypoint": "abc", "runtime_env": {"hi": "hi2"}}, JobSubmitRequest
        )
        assert r.entrypoint == "abc"
        assert r.runtime_env == {"hi": "hi2"}

        with pytest.raises(TypeError, match="must be a dict"):
            validate_request_type(
                {"entrypoint": "abc", "runtime_env": 123}, JobSubmitRequest
            )

        with pytest.raises(TypeError, match="keys must be strings"):
            validate_request_type(
                {"entrypoint": "abc", "runtime_env": {1: "hi"}}, JobSubmitRequest
            )

    def test_validate_metadata(self):
        r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
        assert r.entrypoint == "abc"
        assert r.metadata is None

        r = validate_request_type(
            {"entrypoint": "abc", "metadata": {"hi": "hi2"}}, JobSubmitRequest
        )
        assert r.entrypoint == "abc"
        assert r.metadata == {"hi": "hi2"}

        with pytest.raises(TypeError, match="must be a dict"):
            validate_request_type(
                {"entrypoint": "abc", "metadata": 123}, JobSubmitRequest
            )

        with pytest.raises(TypeError, match="keys must be strings"):
            validate_request_type(
                {"entrypoint": "abc", "metadata": {1: "hi"}}, JobSubmitRequest
            )

        with pytest.raises(TypeError, match="values must be strings"):
            validate_request_type(
                {"entrypoint": "abc", "metadata": {"hi": 1}}, JobSubmitRequest
            )

    def test_validate_entrypoint_label_selector(self):
        r = validate_request_type(
            {
                "entrypoint": "abc",
                "entrypoint_label_selector": {"fragile_node": "!1"},
            },
            JobSubmitRequest,
        )
        assert r.entrypoint_label_selector == {"fragile_node": "!1"}

        with pytest.raises(TypeError, match="must be a dict"):
            validate_request_type(
                {"entrypoint": "abc", "entrypoint_label_selector": "bad"},
                JobSubmitRequest,
            )

        with pytest.raises(TypeError, match="keys must be strings"):
            validate_request_type(
                {"entrypoint": "abc", "entrypoint_label_selector": {1: "bad"}},
                JobSubmitRequest,
            )

        with pytest.raises(TypeError, match="values must be strings"):
            validate_request_type(
                {"entrypoint": "abc", "entrypoint_label_selector": {"k": 1}},
                JobSubmitRequest,
            )

    def test_entrypoint_resources_disallow_strings(self):
        with pytest.raises(TypeError, match="values must be numbers"):
            validate_request_type(
                {"entrypoint": "abc", "entrypoint_resources": {"Custom": "1"}},
                JobSubmitRequest,
            )


def test_uri_to_http_and_back():
    assert uri_to_http_components("gcs://hello.zip") == ("gcs", "hello.zip")
    assert uri_to_http_components("gcs://hello.whl") == ("gcs", "hello.whl")

    with pytest.raises(ValueError, match="'blah' is not a valid Protocol"):
        uri_to_http_components("blah://halb.zip")

    with pytest.raises(ValueError, match="does not end in .zip or .whl"):
        assert uri_to_http_components("gcs://hello.not_zip")

    with pytest.raises(ValueError, match="does not end in .zip or .whl"):
        assert uri_to_http_components("gcs://hello")

    assert http_uri_components_to_uri("gcs", "hello.zip") == "gcs://hello.zip"
    assert http_uri_components_to_uri("blah", "halb.zip") == "blah://halb.zip"
    assert http_uri_components_to_uri("blah", "halb.whl") == "blah://halb.whl"

    for original_uri in ["gcs://hello.zip", "gcs://fasdf.whl"]:
        new_uri = http_uri_components_to_uri(*uri_to_http_components(original_uri))
        assert new_uri == original_uri


def test_dynamic_status_message():
    info = JobInfo(
        status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_num_cpus=1
    )
    assert "may be waiting for resources" in info.message

    info = JobInfo(
        status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_num_gpus=1
    )
    assert "may be waiting for resources" in info.message

    info = JobInfo(status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_memory=4)
    assert "may be waiting for resources" in info.message

    info = JobInfo(
        status=JobStatus.PENDING,
        entrypoint="echo hi",
        entrypoint_resources={"Custom": 1},
    )
    assert "may be waiting for resources" in info.message

    info = JobInfo(
        status=JobStatus.PENDING, entrypoint="echo hi", runtime_env={"conda": "env"}
    )
    assert "may be waiting for the runtime environment" in info.message


def test_job_info_to_json():
    info = JobInfo(
        status=JobStatus.PENDING,
        entrypoint="echo hi",
        entrypoint_num_cpus=1,
        entrypoint_num_gpus=1,
        entrypoint_memory=4,
        entrypoint_resources={"Custom": 1},
        runtime_env={"pip": ["pkg"]},
    )
    expected_items = {
        "status": "PENDING",
        "message": (
            "Job has not started yet. It may be waiting for resources "
            "(CPUs, GPUs, memory, custom resources) to become available. "
            "It may be waiting for the runtime environment to be set up."
        ),
        "entrypoint": "echo hi",
        "entrypoint_num_cpus": 1,
        "entrypoint_num_gpus": 1,
        "entrypoint_memory": 4,
        "entrypoint_resources": {"Custom": 1},
        "runtime_env_json": '{"pip": ["pkg"]}',
    }

    # Check that the expected items are in the JSON.
    assert expected_items.items() <= info.to_json().items()

    new_job_info = JobInfo.from_json(info.to_json())
    assert new_job_info == info

    # If `status` is just a string, then operations like status.is_terminal()
    # would fail, so we should make sure that it's a JobStatus.
    assert isinstance(new_job_info.status, JobStatus)


def test_job_info_json_to_proto():
    """Test that JobInfo JSON can be converted to JobsAPIInfo protobuf."""
    info = JobInfo(
        status=JobStatus.PENDING,
        entrypoint="echo hi",
        error_type=JobErrorType.JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE,
        start_time=123,
        end_time=456,
        metadata={"hi": "hi2"},
        entrypoint_num_cpus=1,
        entrypoint_num_gpus=1,
        entrypoint_memory=4,
        entrypoint_resources={"Custom": 1},
        runtime_env={"pip": ["pkg"]},
        driver_agent_http_address="http://localhost:1234",
        driver_node_id="node_id",
    )
    info_json = json.dumps(info.to_json())
    info_proto = Parse(info_json, JobsAPIInfo())
    assert info_proto.status == "PENDING"
    assert info_proto.entrypoint == "echo hi"
    assert info_proto.start_time == 123
    assert info_proto.end_time == 456
    assert info_proto.metadata == {"hi": "hi2"}
    assert info_proto.entrypoint_num_cpus == 1
    assert info_proto.entrypoint_num_gpus == 1
    assert info_proto.entrypoint_memory == 4
    assert info_proto.entrypoint_resources == {"Custom": 1}
    assert info_proto.runtime_env_json == '{"pip": ["pkg"]}'
    assert info_proto.message == (
        "Job has not started yet. It may be waiting for resources "
        "(CPUs, GPUs, memory, custom resources) to become available. "
        "It may be waiting for the runtime environment to be set up."
    )
    assert info_proto.error_type == "JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE"
    assert info_proto.driver_agent_http_address == "http://localhost:1234"
    assert info_proto.driver_node_id == "node_id"

    minimal_info = JobInfo(status=JobStatus.PENDING, entrypoint="echo hi")
    minimal_info_json = json.dumps(minimal_info.to_json())
    minimal_info_proto = Parse(minimal_info_json, JobsAPIInfo())
    assert minimal_info_proto.status == "PENDING"
    assert minimal_info_proto.entrypoint == "echo hi"
    for unset_optional_field in [
        "entrypoint_num_cpus",
        "entrypoint_num_gpus",
        "entrypoint_memory",
        "runtime_env_json",
        "error_type",
        "driver_agent_http_address",
        "driver_node_id",
    ]:
        assert not minimal_info_proto.HasField(unset_optional_field)


def test_get_all_jobs_filters_out_none_job_info():
    prefix = JobInfoStorageClient.JOB_DATA_KEY_PREFIX
    mock_gcs = MagicMock()
    mock_gcs.async_internal_kv_keys = AsyncMock(
        return_value=[
            (prefix + "job1").encode(),
            (prefix + "job2").encode(),
        ]
    )

    storage = JobInfoStorageClient(mock_gcs)
    job_info_1 = JobInfo(status=JobStatus.RUNNING, entrypoint="echo 1")

    async def mock_get_info(job_id, timeout=30):
        if job_id == "job1":
            return job_info_1
        return None

    storage.get_info = mock_get_info

    result = asyncio.run(storage.get_all_jobs())

    assert result == {"job1": job_info_1}
    for job_id, job_info in result.items():
        asdict(job_info)  # This should not raise an exception


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))