import os
import time
import uuid
from pathlib import Path
from typing import Dict, List, Optional
from unittest.mock import MagicMock

import ray
from ray.train import BackendConfig, Checkpoint
from ray.train._internal.data_config import DataConfig
from ray.train.backend import Backend
from ray.train.context import TrainContext
from ray.train.v2._internal.execution.context import (
    DistributedContext,
    TrainRunContext,
)
from ray.train.v2._internal.execution.failure_handling import (
    FailureDecision,
    FailurePolicy,
)
from ray.train.v2._internal.execution.scaling_policy import (
    NoopDecision,
    ScalingDecision,
    ScalingPolicy,
)
from ray.train.v2._internal.execution.storage import StorageContext
from ray.train.v2._internal.execution.training_report import _TrainingReport
from ray.train.v2._internal.execution.worker_group import (
    WorkerGroup,
    WorkerGroupContext,
    WorkerGroupPollStatus,
    WorkerGroupState,
    WorkerStatus,
)
from ray.train.v2._internal.execution.worker_group.execution_group import ReplicaGroup
from ray.train.v2._internal.state.schema import (
    ActorStatus,
    BackendConfig as BackendConfigSchema,
    CheckpointConfig as CheckpointConfigSchema,
    DataConfig as DataConfigSchema,
    DataExecutionOptions,
    FailureConfig as FailureConfigSchema,
    RunAttemptStatus,
    RunConfig as RunConfigSchema,
    RunSettings,
    RunStatus,
    ScalingConfig as ScalingConfigSchema,
    TrainResources,
    TrainRun,
    TrainRunAttempt,
    TrainWorker,
)
from ray.train.v2._internal.state.util import execution_options_to_model
from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
from ray.train.v2.api.exceptions import TrainingFailedError
from ray.train.v2.api.validation_config import ValidationTaskConfig


class MockReplicaGroupBackend(Backend):
    has_replica_groups = True


class MockReplicaGroupBackendConfig(BackendConfig):
    @property
    def backend_cls(self):
        return MockReplicaGroupBackend


@ray.remote
class Counter:
    def __init__(self):
        self.count = 0

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

    def get_count(self):
        return self.count


class DummyWorkerGroup(WorkerGroup):

    _start_failure = None
    _poll_failure = None

    # TODO: Clean this up and use Mocks instead.
    def __init__(
        self,
        train_run_context: TrainRunContext,
        worker_group_context: WorkerGroupContext,
        callbacks=None,
    ):
        self._num_workers = worker_group_context.num_workers
        self._worker_group_state = None
        self._worker_statuses = {}
        self._replaced_replica_groups: List[int] = []
        self._replica_groups = None
        self._latest_poll_status: Optional[WorkerGroupPollStatus] = None

    def poll_status(self, *args, **kwargs) -> WorkerGroupPollStatus:
        if self._poll_failure:
            raise self._poll_failure
        return WorkerGroupPollStatus(
            worker_statuses=self._worker_statuses,
        )

    def _start(self):
        num_workers = self._num_workers
        if self._start_failure:
            raise self._start_failure

        workers = [MagicMock() for i in range(num_workers)]
        self._worker_group_state = WorkerGroupState(
            start_time=time_monotonic(),
            workers=workers,
            placement_group_handle=MagicMock(),
            sync_actor=None,
        )

        self._worker_statuses = {
            i: WorkerStatus(running=True, error=None) for i in range(num_workers)
        }

        self._replica_groups = [
            ReplicaGroup([workers[i]], resources_per_worker={})
            for i in range(num_workers)
        ]

    def shutdown(self):
        self._worker_group_state = None

    def abort(self):
        pass

    def replace_replica_group(self, replica_group_index: int):
        self._replaced_replica_groups.append(replica_group_index)

    # === Test methods ===
    def clear_worker(self):
        for worker_status in self._worker_statuses.values():
            worker_status.error = None
            worker_status.running = True

    def error_worker(self, worker_index):
        status = self._worker_statuses[worker_index]
        status.error = RuntimeError(f"Worker {worker_index} failed")

    def finish_worker(self, worker_index):
        status = self._worker_statuses[worker_index]
        status.running = False

    @classmethod
    def set_start_failure(cls, start_failure):
        cls._start_failure = start_failure

    @classmethod
    def set_poll_failure(cls, poll_failure):
        cls._poll_failure = poll_failure


class MockScalingPolicy(ScalingPolicy):
    def __init__(self, scaling_config):
        self._recovery_decision_queue = []
        self._monitor_decision_queue = []

        super().__init__(scaling_config)

    def _get_num_workers_for_resource_request(self) -> int:
        return self.scaling_config.num_workers

    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
        if self._recovery_decision_queue:
            return self._recovery_decision_queue.pop(0)
        return NoopDecision()

    def make_decision_for_running_worker_group(
        self,
        worker_group_state: WorkerGroupState,
        worker_group_status: WorkerGroupPollStatus,
    ) -> ScalingDecision:
        if self._monitor_decision_queue:
            return self._monitor_decision_queue.pop(0)
        return NoopDecision()

    # === Test methods ===
    def queue_recovery_decision(self, decision):
        self._recovery_decision_queue.append(decision)

    def queue_monitor_decision(self, decision):
        self._monitor_decision_queue.append(decision)


class MockFailurePolicy(FailurePolicy):
    def __init__(self, failure_config):
        self._decision_queue = []

        super().__init__(failure_config)

    def make_decision(
        self, training_failed_error: TrainingFailedError
    ) -> FailureDecision:
        if self._decision_queue:
            return self._decision_queue.pop(0)
        return FailureDecision.NOOP

    # === Test methods ===
    def queue_decision(self, decision):
        self._decision_queue.append(decision)


class DummyObjectRefWrapper(ObjectRefWrapper):
    """Mock object that returns the object passed in without going through ray.put."""

    def __init__(self, obj):
        self._obj = obj

    def get(self):
        return self._obj


_RUN_ID = "mock_run_id"


def create_mock_train_run(
    status: RunStatus = RunStatus.RUNNING,
    controller_actor_id: Optional[str] = None,
    end_time_ns: Optional[int] = None,
    id: Optional[str] = None,
    status_detail: Optional[str] = None,
    train_loop_config: Optional[Dict] = None,
):
    return TrainRun(
        schema_version=0,
        id=id or _RUN_ID,
        name="test_run",
        job_id=uuid.uuid4().hex,
        controller_actor_id=controller_actor_id or uuid.uuid4().hex,
        status=status,
        status_detail=status_detail,
        start_time_ns=time.time_ns(),
        end_time_ns=end_time_ns,
        controller_log_file_path="/tmp/ray/session_xxx/logs/train/ray-train-app-controller.log",
        framework_versions={"ray": ray.__version__},
        run_settings=RunSettings(
            train_loop_config=train_loop_config,
            backend_config=BackendConfigSchema(framework=None, config={}),
            scaling_config=ScalingConfigSchema(
                num_workers=1,
                use_gpu=False,
                resources_per_worker=None,
                placement_strategy="PACK",
                accelerator_type=None,
                use_tpu=False,
                topology=None,
                bundle_label_selector=None,
            ),
            datasets=["dataset_1"],
            data_config=DataConfigSchema(
                datasets_to_split="all",
                data_execution_options=DataExecutionOptions(
                    default=execution_options_to_model(
                        DataConfig.default_ingest_options()
                    ),
                ),
                enable_shard_locality=True,
            ),
            run_config=RunConfigSchema(
                name="test_run",
                failure_config=FailureConfigSchema(
                    max_failures=0, controller_failure_limit=-1
                ),
                worker_runtime_env={"type": "conda"},
                checkpoint_config=CheckpointConfigSchema(
                    num_to_keep=None,
                    checkpoint_score_attribute=None,
                    checkpoint_score_order="max",
                ),
                storage_path="s3://bucket/path",
                storage_filesystem=None,
            ),
        ),
    )


def create_mock_train_run_attempt(
    attempt_id: str = "mock_attempt_id",
    status: RunAttemptStatus = RunAttemptStatus.RUNNING,
    end_time_ns: Optional[int] = None,
    run_id: Optional[str] = None,
    worker_status: Optional[ActorStatus] = ActorStatus.ALIVE,
    status_detail: Optional[str] = None,
):
    worker = TrainWorker(
        world_rank=0,
        local_rank=0,
        node_rank=0,
        actor_id=uuid.uuid4().hex,
        node_id=uuid.uuid4().hex,
        node_ip="127.0.0.1",
        pid=1234,
        gpu_ids=[0],
        status=worker_status,
        resources=TrainResources(resources={"CPU": 1}),
        log_file_path="/tmp/ray/session_xxx/logs/train/ray-train-app-worker.log",
    )

    return TrainRunAttempt(
        schema_version=0,
        attempt_id=attempt_id,
        run_id=run_id or _RUN_ID,
        status=status,
        status_detail=status_detail,
        start_time_ns=time.time_ns(),
        resources=[TrainResources(resources={"CPU": 1})],
        workers=[worker],
        end_time_ns=end_time_ns,
    )


def create_dummy_run_context(**kwargs: dict) -> TrainRunContext:
    """Create a standardized TrainRunContext for testing.

    Args:
        **kwargs: Optional overrides for the default configuration.

    Returns:
        TrainRunContext: A standardized TrainRunContext instance for testing.
    """
    from ray.train import BackendConfig, DataConfig
    from ray.train.v2._internal.execution.context import TrainRunContext
    from ray.train.v2.api.config import RunConfig, ScalingConfig

    config = dict(
        run_config=RunConfig(name="test"),
        train_loop_config={},
        scaling_config=ScalingConfig(num_workers=1),
        backend_config=BackendConfig(),
        dataset_config=DataConfig(),
    )
    config.update(kwargs)
    return TrainRunContext(**config)


class DummyTrainContext(TrainContext):
    """A dummy TrainContext subclass for testing."""

    def __init__(self):
        self.train_run_context = create_dummy_run_context()
        self.distributed_context = DistributedContext(
            world_rank=0,
            world_size=1,
            local_rank=0,
            local_world_size=1,
            node_rank=0,
        )
        # Mock everything else since we don't need the actual functionality
        self.execution_context = MagicMock()
        self.storage_context = MagicMock()
        self.dataset_shards = {}

    def get_run_config(self):
        return self.train_run_context.run_config


def create_dummy_train_context() -> TrainContext:
    """Create a standardized TrainContext for testing.

    Returns:
        TrainContext: A standardized TrainContext instance for testing.
    """
    return DummyTrainContext()


def create_dummy_training_reports(
    num_results: int,
    storage_context: StorageContext,
    include_metrics: bool = True,
    include_validation: bool = False,
    starting_checkpoint_index: int = 0,
) -> List[_TrainingReport]:
    training_results = []
    for i in range(num_results):
        metrics = {"score": i} if include_metrics else {}
        validation = (
            ValidationTaskConfig(fn_kwargs={"arg": i}) if include_validation else False
        )
        checkpoint_path = os.path.join(
            storage_context.experiment_fs_path,
            f"checkpoint_{starting_checkpoint_index + i}",
        )
        os.makedirs(checkpoint_path, exist_ok=True)
        training_results.append(
            _TrainingReport(
                checkpoint=Checkpoint(
                    path=Path(checkpoint_path).as_posix(),
                    filesystem=storage_context.storage_filesystem,
                ),
                metrics=metrics,
                validation=validation,
            )
        )
    return training_results