import uuid
from typing import Optional
from unittest.mock import create_autospec

import pytest

import ray
from ray.train import CheckpointConfig
from ray.train._internal.session import _TrainingResult
from ray.train.v2._internal.exceptions import CheckpointManagerInitializationError
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
    CheckpointManager,
)
from ray.train.v2._internal.execution.storage import StorageContext
from ray.train.v2._internal.execution.worker_group import Worker
from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
from ray.train.v2.tests.util import create_dummy_training_reports


@pytest.fixture(autouse=True, scope="module")
def ray_start_4_cpus():
    ray.init(num_cpus=4)
    yield
    ray.shutdown()


def _checkpoint_managers_equal(cm1: CheckpointManager, cm2: CheckpointManager) -> bool:
    """
    Compare two checkpoint managers for equality.
    Ignore uuid differences of all the checkpoints recorded.
    """

    def _training_results_equal(
        tr1: Optional[_TrainingResult], tr2: Optional[_TrainingResult]
    ) -> bool:
        if not tr1 and not tr2:
            return True
        if not tr1 or not tr2:
            return False
        return (
            tr1.metrics == tr2.metrics
            and tr1.checkpoint.path == tr2.checkpoint.path
            and tr1.checkpoint.filesystem == tr2.checkpoint.filesystem
        )

    def _checkpoint_to_report_indices_equal(
        cm1: CheckpointManager, cm2: CheckpointManager
    ) -> bool:
        # Do this because Checkpoint and Filesystem are not hashable.
        cm1_dict = {
            checkpoint.path: report_index
            for checkpoint, report_index in cm1._checkpoint_to_report_index.items()
        }
        cm2_dict = {
            checkpoint.path: report_index
            for checkpoint, report_index in cm2._checkpoint_to_report_index.items()
        }
        return cm1_dict == cm2_dict

    if cm1._checkpoint_config != cm2._checkpoint_config:
        return False
    if not _training_results_equal(
        cm1.latest_checkpoint_result, cm2.latest_checkpoint_result
    ):
        return False
    if not _training_results_equal(
        cm1.best_checkpoint_result, cm2.best_checkpoint_result
    ):
        return False
    if len(cm1.best_checkpoint_results) != len(cm2.best_checkpoint_results):
        return False
    for tr1, tr2 in zip(cm1.best_checkpoint_results, cm2.best_checkpoint_results):
        if not _training_results_equal(tr1, tr2):
            return False
    if cm1._current_report_index != cm2._current_report_index:
        return False
    if not _checkpoint_to_report_indices_equal(cm1, cm2):
        return False
    return True


@pytest.mark.parametrize(
    "checkpoint_config",
    [
        CheckpointConfig(),
        CheckpointConfig(
            num_to_keep=1,
            checkpoint_score_attribute="score",
            checkpoint_score_order="max",
        ),
    ],
)
@pytest.mark.asyncio
async def test_save_load_state_equivalence(
    monkeypatch, tmp_path, checkpoint_config: CheckpointConfig
):
    # Use async here because register_checkpoint creates an async task

    # Mock the delete function as we don't want report checkpoints to be deleted.
    monkeypatch.setattr(
        ray.train.v2._internal.execution.checkpoint.checkpoint_manager,
        "delete_fs_path",
        lambda *args, **kwargs: None,
    )
    exp_name = f"checkpoint_manager_test-{uuid.uuid4().hex}"

    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name=exp_name,
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=checkpoint_config,
    )
    training_reports = create_dummy_training_reports(
        num_results=2, storage_context=storage_context
    ) + create_dummy_training_reports(
        num_results=1,
        storage_context=storage_context,
        include_validation=True,
        starting_checkpoint_index=2,
    )

    # Register the training results into checkpoint manager
    for i, tr in enumerate(training_reports):
        checkpoint_manager.register_checkpoint(tr)
        assert checkpoint_manager._current_report_index == i + 1
        loaded_checkpoint_manager = CheckpointManager(
            storage_context=storage_context,
            checkpoint_config=checkpoint_config,
        )
        assert _checkpoint_managers_equal(checkpoint_manager, loaded_checkpoint_manager)


@pytest.mark.parametrize(
    "json_state,match",
    [
        (
            '{"dummy": "1", "dummy_dict": {"key": "value"}}',
            "You are loading a checkpoint manager snapshot saved with an unknown Ray version but",
        ),
        ('{"ray_version": "2.0.0", "dummy": "1", "dummy_dict": {"key": "value"', None),
        (
            '{"ray_version": "2.0.0", "dummy": "1", "dummy_dict": {"key": "value"}}',
            "You are loading a checkpoint manager snapshot saved with Ray version 2.0.0 but",
        ),
    ],
)
def test_load_state_error(tmp_path, json_state, match):

    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name="load_state_error_experiment",
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=CheckpointConfig(),
    )
    with pytest.raises(
        CheckpointManagerInitializationError,
        match=match,
    ):
        checkpoint_manager._load_state(json_state)


@pytest.mark.asyncio
async def test_before_init_train_context(tmp_path):

    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name="my_experiment_name",
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=CheckpointConfig(),
    )
    workers = [create_autospec(Worker, instance=True) for _ in range(4)]

    # Assert without a checkpoint.
    assert checkpoint_manager.before_init_train_context(workers) == {
        "checkpoint": [None] * 4,
        "current_report_index": [0] * 4,
    }

    # Assert with a checkpoint
    latest_checkpoint_report = create_dummy_training_reports(1, storage_context)[0]
    checkpoint_manager.register_checkpoint(latest_checkpoint_report)
    assert checkpoint_manager.before_init_train_context(workers) == {
        "checkpoint": [latest_checkpoint_report.checkpoint] * 4,
        "current_report_index": [1] * 4,
    }


@pytest.mark.asyncio
async def test_pending_checkpoint_management(tmp_path):
    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name="pending_checkpoint_management_experiment",
    )
    checkpoint_config = CheckpointConfig(
        num_to_keep=1,
        checkpoint_score_attribute="score",
        checkpoint_score_order="max",
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=checkpoint_config,
    )
    (
        low_initial_high_final_training_report,
        high_initial_low_final_training_report,
        final_training_report,
    ) = create_dummy_training_reports(
        num_results=3, storage_context=storage_context, include_validation=True
    )
    final_training_report.validation = False
    scoreless_training_report = create_dummy_training_reports(
        num_results=1,
        storage_context=storage_context,
        include_metrics=False,
        starting_checkpoint_index=3,
    )[0]

    # Register pending/final/unknown checkpoints and verify their storage
    checkpoint_manager.register_checkpoint(low_initial_high_final_training_report)
    checkpoint_manager.register_checkpoint(final_training_report)
    checkpoint_manager.register_checkpoint(scoreless_training_report)
    checkpoint_manager.register_checkpoint(high_initial_low_final_training_report)
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        low_initial_high_final_training_report.checkpoint,  # keep pending
        high_initial_low_final_training_report.checkpoint,  # keep pending/latest
        final_training_report.checkpoint,  # keep highest final score so far
    ]

    # Assert checkpoint state after all tasks are done
    checkpoint_manager.update_checkpoints_with_validation_result(
        {
            low_initial_high_final_training_report.checkpoint: (
                {"score": 200},
                ReportedCheckpointStatus.VALIDATED,
            ),
            high_initial_low_final_training_report.checkpoint: (
                {"score": 100},
                ReportedCheckpointStatus.VALIDATED,
            ),
        }
    )
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        high_initial_low_final_training_report.checkpoint,  # keep latest checkpoint
        low_initial_high_final_training_report.checkpoint,  # keep highest score checkpoint
    ]


@pytest.mark.asyncio
async def test_pending_checkpoint_management_break_ties_by_report_index(tmp_path):
    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name="pending_checkpoint_management_break_ties_by_report_index_experiment",
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=CheckpointConfig(),
    )
    training_reports = create_dummy_training_reports(
        num_results=2, storage_context=storage_context, include_validation=True
    )
    checkpoint_manager.register_checkpoint(training_reports[0])
    checkpoint_manager.register_checkpoint(training_reports[1])
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        training_reports[0].checkpoint,
        training_reports[1].checkpoint,
    ]
    checkpoint_manager.update_checkpoints_with_validation_result(
        {
            training_reports[1].checkpoint: ({}, ReportedCheckpointStatus.VALIDATED),
        }
    )
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        training_reports[0].checkpoint,
        training_reports[1].checkpoint,
    ]
    checkpoint_manager.update_checkpoints_with_validation_result(
        {
            training_reports[0].checkpoint: ({}, ReportedCheckpointStatus.VALIDATED),
        }
    )
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        training_reports[0].checkpoint,
        training_reports[1].checkpoint,
    ]


@pytest.mark.asyncio
async def test_pending_checkpoint_management_finalized_checkpoint(tmp_path):
    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name="pending_checkpoint_management_experiment",
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=CheckpointConfig(
            checkpoint_score_attribute="score",
            checkpoint_score_order="max",
        ),
    )
    training_reports = create_dummy_training_reports(
        num_results=2, storage_context=storage_context
    )
    checkpoint_manager.register_checkpoint(training_reports[0])
    checkpoint_manager.register_checkpoint(training_reports[1])
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        training_reports[0].checkpoint,
        training_reports[1].checkpoint,
    ]
    checkpoint_manager.update_checkpoints_with_validation_result(
        {
            training_reports[0].checkpoint: (
                {"score": 100},
                ReportedCheckpointStatus.VALIDATED,
            ),
        }
    )
    assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
        training_reports[0].checkpoint,
        training_reports[1].checkpoint,
    ]


def test_update_checkpoints_with_metrics_not_in_checkpoint_results(tmp_path):
    storage_context = StorageContext(
        storage_path=tmp_path,
        experiment_dir_name="update_checkpoints_with_metrics_error_experiment",
    )
    checkpoint_manager = CheckpointManager(
        storage_context=storage_context,
        checkpoint_config=CheckpointConfig(),
    )
    training_reports = create_dummy_training_reports(
        num_results=1, storage_context=storage_context
    )
    checkpoint_manager._pending_training_results[training_reports[0].checkpoint] = (
        _TrainingResult(training_reports[0].checkpoint, training_reports[0].metrics),
        training_reports[0].validation,
    )
    with pytest.raises(ValueError):
        checkpoint_manager.update_checkpoints_with_validation_result(
            {
                training_reports[0].checkpoint: (
                    {"score": 100},
                    ReportedCheckpointStatus.VALIDATED,
                )
            }
        )


if __name__ == "__main__":
    import sys

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