"""Integration tests for token-based authentication in Ray (part 1).

Cluster lifecycle, connection, and end-to-end operation tests.
"""

import os
import subprocess
import sys
from pathlib import Path
from typing import Optional

import pytest

import ray
from ray._common.test_utils import wait_for_condition
from ray._private.test_utils import client_test_enabled

try:
    from ray._raylet import AuthenticationTokenLoader

    _RAYLET_AVAILABLE = True
except ImportError:
    _RAYLET_AVAILABLE = False
    AuthenticationTokenLoader = None

from ray._private.authentication_test_utils import (
    clear_auth_token_sources,
    reset_auth_token_state,
    set_auth_mode,
    set_env_auth_token,
)

pytestmark = pytest.mark.skipif(
    not _RAYLET_AVAILABLE,
    reason="Authentication tests require ray._raylet (not available in minimal installs)",
)


def _run_ray_start_and_verify_status(
    args: list, env: dict, expect_success: bool = True, timeout: int = 30
) -> subprocess.CompletedProcess:
    """Helper to run ray start command with proper error handling."""
    result = subprocess.run(
        ["ray", "start"] + args,
        env={"RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER": "1", **env},
        capture_output=True,
        text=True,
        timeout=timeout,
    )

    if expect_success:
        assert result.returncode == 0, (
            f"ray start should have succeeded. "
            f"stdout: {result.stdout}, stderr: {result.stderr}"
        )
    else:
        assert result.returncode != 0, (
            f"ray start should have failed but succeeded. "
            f"stdout: {result.stdout}, stderr: {result.stderr}"
        )
        # Check that error message mentions token
        error_output = result.stdout + result.stderr
        assert (
            "authentication token" in error_output.lower()
            or "token" in error_output.lower()
        ), f"Error message should mention token. Got: {error_output}"

    return result


def _cleanup_ray_start(env: Optional[dict] = None):
    """Helper to clean up ray start processes."""
    # Ensure any ray.init() connection is closed first
    if ray.is_initialized():
        ray.shutdown()

    # Stop with a longer timeout
    subprocess.run(
        ["ray", "stop", "--force"],
        env=env,
        capture_output=True,
        timeout=60,  # Increased timeout for flaky cleanup
        check=False,  # Don't raise on non-zero exit
    )

    # Wait for ray processes to actually stop
    def ray_stopped():
        result = subprocess.run(
            ["ray", "status"],
            capture_output=True,
            check=False,
        )
        # ray status returns non-zero when no cluster is running
        return result.returncode != 0

    try:
        wait_for_condition(ray_stopped, timeout=10, retry_interval_ms=500)
    except Exception:
        # Best effort - don't fail the test if we can't verify it stopped
        pass


@pytest.fixture(autouse=True)
def _auto_clean_token_sources(clean_token_sources):
    """Opt in to the shared clean_token_sources fixture for every test."""
    yield


@pytest.mark.skipif(
    client_test_enabled(),
    reason="This test is for starting a new local cluster, not compatible with client mode",
)
def test_local_cluster_generates_token():
    """Test ray.init() generates token for local cluster when auth_mode=token is set."""
    # Ensure no token exists
    default_token_path = Path.home() / ".ray" / "auth_token"
    assert (
        not default_token_path.exists()
    ), f"Token file already exists at {default_token_path}"

    # Enable token auth via environment variable
    set_auth_mode("token")
    reset_auth_token_state()

    # Initialize Ray with token auth
    ray.init()

    try:
        # Verify token file was created
        assert default_token_path.exists(), (
            f"Token file was not created at {default_token_path}. "
            f"HOME={os.environ.get('HOME')}, "
            f"Files in {default_token_path.parent}: {list(default_token_path.parent.iterdir()) if default_token_path.parent.exists() else 'directory does not exist'}"
        )
        token = default_token_path.read_text().strip()
        assert len(token) == 64
        assert all(c in "0123456789abcdef" for c in token)

        # Verify cluster is working
        assert ray.is_initialized()

    finally:
        ray.shutdown()


def test_connect_without_token_raises_error(setup_cluster_with_token_auth):
    """Test ray.init(address=...) without token fails when auth_mode=token is set."""
    cluster_info = setup_cluster_with_token_auth
    cluster = cluster_info["cluster"]

    # Disconnect the current driver session and drop token state before retrying.
    ray.shutdown()
    set_auth_mode("disabled")
    clear_auth_token_sources(remove_default=True)
    reset_auth_token_state()

    # Ensure no token exists
    token_loader = AuthenticationTokenLoader.instance()
    assert not token_loader.has_token()

    # Try to connect to the cluster without a token - should raise RuntimeError
    with pytest.raises(ConnectionError):
        ray.init(address=cluster.address)


@pytest.mark.parametrize(
    "token,expected_status",
    [
        (None, 401),  # No token -> Unauthorized
        ("wrong_token", 403),  # Wrong token -> Forbidden
    ],
    ids=["no_token", "wrong_token"],
)
def test_state_api_auth_failure(token, expected_status, setup_cluster_with_token_auth):
    """Test that state API calls fail with missing or incorrect token."""
    import requests

    cluster_info = setup_cluster_with_token_auth
    dashboard_url = cluster_info["dashboard_url"]

    # Make direct HTTP request to state API endpoint
    headers = {}
    if token is not None:
        headers["Authorization"] = f"Bearer {token}"

    response = requests.get(f"{dashboard_url}/api/v0/actors", headers=headers)

    assert response.status_code == expected_status, (
        f"State API should return {expected_status}, got {response.status_code}: "
        f"{response.text}"
    )


@pytest.mark.parametrize("tokens_match", [True, False])
def test_cluster_token_authentication(tokens_match, setup_cluster_with_token_auth):
    """Test cluster authentication with matching and non-matching tokens."""
    cluster_info = setup_cluster_with_token_auth
    cluster = cluster_info["cluster"]
    cluster_token = cluster_info["token"]

    # Reconfigure the driver token state to simulate fresh connections.
    ray.shutdown()
    set_auth_mode("token")

    if tokens_match:
        client_token = cluster_token  # Same token - should succeed
    else:
        client_token = "b" * 64  # Different token - should fail

    set_env_auth_token(client_token)
    reset_auth_token_state()

    if tokens_match:
        # Should succeed - test gRPC calls work
        ray.init(address=cluster.address)

        obj_ref = ray.put("test_data")
        result = ray.get(obj_ref)
        assert result == "test_data"

        @ray.remote
        def test_func():
            return "success"

        result = ray.get(test_func.remote())
        assert result == "success"

        ray.shutdown()

    else:
        # Should fail - connection or gRPC calls should fail
        with pytest.raises((ConnectionError, RuntimeError)):
            ray.init(address=cluster.address)
            try:
                ray.put("test")
            finally:
                ray.shutdown()


@pytest.mark.skipif(
    client_test_enabled(),
    reason="Uses subprocess ray start, not compatible with client mode",
)
@pytest.mark.parametrize("is_head", [True, False])
def test_ray_start_without_token_raises_error(is_head, request):
    """Test that ray start fails when auth_mode=token but no token exists."""
    # Set up environment with token auth enabled but no token
    env = os.environ.copy()
    env["RAY_AUTH_MODE"] = "token"
    env.pop("RAY_AUTH_TOKEN", None)
    env.pop("RAY_AUTH_TOKEN_PATH", None)

    # Ensure no default token file exists (already cleaned by fixture)
    default_token_path = Path.home() / ".ray" / "auth_token"
    assert not default_token_path.exists()

    # When specifying an address, we need a head node to connect to
    cluster_info = None
    if not is_head:
        cluster_info = request.getfixturevalue("setup_cluster_with_token_auth")
        cluster = cluster_info["cluster"]
        ray.shutdown()

    # Prepare arguments
    if is_head:
        args = ["--head", "--port=0"]
    else:
        args = [f"--address={cluster.address}"]

    # Try to start node - should fail
    _run_ray_start_and_verify_status(args, env, expect_success=False)


@pytest.mark.skipif(
    client_test_enabled(),
    reason="Uses subprocess ray start, not compatible with client mode",
)
def test_ray_start_head_with_token_succeeds():
    """Test that ray start --head succeeds when token auth is enabled with a valid token."""
    # Set up environment with token auth and a valid token
    test_token = "a" * 64
    env = os.environ.copy()
    env["RAY_AUTH_TOKEN"] = test_token
    env["RAY_AUTH_MODE"] = "token"

    try:
        # Start head node - should succeed
        _run_ray_start_and_verify_status(
            ["--head", "--port=0"], env, expect_success=True
        )

        # Verify we can connect to the cluster with ray.init()
        set_env_auth_token(test_token)
        set_auth_mode("token")
        reset_auth_token_state()

        # Wait for cluster to be ready
        def cluster_ready():
            try:
                ray.init(address="auto")
                return True
            except Exception:
                return False

        wait_for_condition(cluster_ready, timeout=10)
        assert ray.is_initialized()

        # Test basic operations work
        @ray.remote
        def test_func():
            return "success"

        result = ray.get(test_func.remote())
        assert result == "success"

    finally:
        # Cleanup handles ray.shutdown() internally
        _cleanup_ray_start(env)


@pytest.mark.skipif(
    client_test_enabled(),
    reason="Uses subprocess ray start, not compatible with client mode",
)
@pytest.mark.parametrize("token_match", ["correct", "incorrect"])
def test_ray_start_address_with_token(token_match, setup_cluster_with_token_auth):
    """Test ray start --address=... with correct or incorrect token."""
    cluster_info = setup_cluster_with_token_auth
    cluster = cluster_info["cluster"]
    cluster_token = cluster_info["token"]

    # Reset the driver connection to reuse the fixture-backed cluster.
    ray.shutdown()
    set_auth_mode("token")

    # Set up environment for worker
    env = os.environ.copy()
    env["RAY_AUTH_MODE"] = "token"

    if token_match == "correct":
        env["RAY_AUTH_TOKEN"] = cluster_token
        expect_success = True
    else:
        env["RAY_AUTH_TOKEN"] = "b" * 64
        expect_success = False

    # Start worker node
    _run_ray_start_and_verify_status(
        [f"--address={cluster.address}", "--num-cpus=1"],
        env,
        expect_success=expect_success,
    )

    if token_match == "correct":
        try:
            # Connect and verify the cluster has 2 nodes (head + worker)
            set_env_auth_token(cluster_token)
            reset_auth_token_state()
            ray.init(address=cluster.address)

            def worker_joined():
                return len(ray.nodes()) >= 2

            wait_for_condition(worker_joined, timeout=10)

            nodes = ray.nodes()
            assert (
                len(nodes) >= 2
            ), f"Expected at least 2 nodes, got {len(nodes)}: {nodes}"

        finally:
            if ray.is_initialized():
                ray.shutdown()
            _cleanup_ray_start(env)


def test_e2e_operations_with_token_auth(setup_cluster_with_token_auth):
    """Test that e2e operations work with token authentication enabled.

    This verifies that with token auth enabled:
    1. Tasks execute successfully
    2. Actors can be created and called
    3. State API works (list_nodes, list_actors, list_tasks)
    4. Job submission works
    """
    cluster_info = setup_cluster_with_token_auth

    # Test 1: Submit a simple task
    @ray.remote
    def simple_task(x):
        return x + 1

    result = ray.get(simple_task.remote(41))
    assert result == 42, f"Task should return 42, got {result}"

    # Test 2: Create and use an actor
    @ray.remote
    class SimpleActor:
        def __init__(self):
            self.value = 0

        def increment(self):
            self.value += 1
            return self.value

    actor = SimpleActor.remote()
    result = ray.get(actor.increment.remote())
    assert result == 1, f"Actor method should return 1, got {result}"

    # Test 3: State API operations (uses HTTP with auth headers)
    from ray.util.state import list_actors, list_nodes, list_tasks

    # List nodes - should include at least the head node
    wait_for_condition(lambda: len(list_nodes()) >= 1)

    # List actors - should include our SimpleActor
    def check_actors():
        actors = list_actors()
        if len(actors) < 1:
            return False
        return "SimpleActor" in actors[0].class_name

    wait_for_condition(check_actors)

    # List tasks - should include completed tasks
    wait_for_condition(lambda: len(list_tasks()) >= 1)

    # Test 4: Submit a job and wait for completion
    from ray.job_submission import JobSubmissionClient

    # Create job submission client (uses HTTP with auth headers)
    client = JobSubmissionClient(address=cluster_info["dashboard_url"])

    # Submit a simple job
    job_id = client.submit_job(
        entrypoint="echo 'Hello from job'",
    )

    # Wait for job to complete
    def job_finished():
        status = client.get_job_status(job_id)
        return status in ["SUCCEEDED", "FAILED", "STOPPED"]

    wait_for_condition(job_finished, timeout=30)

    final_status = client.get_job_status(job_id)
    assert (
        final_status == "SUCCEEDED"
    ), f"Job should succeed, got status: {final_status}"


def test_logs_api_with_token_auth(setup_cluster_with_token_auth):
    """Test that log APIs work with token authentication enabled."""
    from ray.util.state import get_log, list_logs

    # Get node ID for log queries
    node_id = ray.nodes()[0]["NodeID"]

    # Test list_logs() with valid auth
    logs = list_logs(node_id=node_id)
    assert isinstance(logs, dict), f"list_logs should return a dict, got {type(logs)}"

    # Test get_log() with valid auth (fetch raylet.out which will always exist)
    chunks_received = 0
    for chunk in get_log(filename="raylet.out", node_id=node_id, tail=10):
        assert isinstance(chunk, str), f"get_log chunk should be str, got {type(chunk)}"
        chunks_received += 1
        break

    assert chunks_received > 0, "Should have received at least one log chunk"


if __name__ == "__main__":
    sys.exit(pytest.main(["-vv", __file__]))