import logging
from collections import defaultdict
from typing import Dict, List, Optional

import ray
from ray.actor import ActorHandle
from ray.train import BackendConfig
from ray.train._internal.data_config import DataConfig
from ray.train.v2._internal.execution.context import DistributedContext
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
    ResizeDecision,
)
from ray.train.v2._internal.execution.worker_group import ActorMetadata, Worker
from ray.train.v2._internal.state.schema import (
    ActorStatus,
    BackendConfig as BackendConfigSchema,
    CheckpointConfig as CheckpointConfigSchema,
    FailureConfig as FailureConfigSchema,
    RunAttemptStatus,
    RunConfig as RunConfigSchema,
    RunSettings,
    RunStatus,
    ScalingConfig as ScalingConfigSchema,
    TrainResources,
    TrainRun,
    TrainRunAttempt,
    TrainWorker,
)
from ray.train.v2._internal.state.state_actor import get_or_create_state_actor
from ray.train.v2._internal.state.util import (
    construct_data_config,
    current_time_ns,
    mark_workers_dead,
    update_train_run_aborted,
    update_train_run_attempt_aborted,
)
from ray.train.v2._internal.util import TrainingFramework
from ray.train.v2.api.config import RunConfig, ScalingConfig

logger = logging.getLogger(__name__)


class TrainStateManager:
    """Manages the state of a train run and run attempts."""

    def __init__(self) -> None:
        self._state_actor = get_or_create_state_actor()
        # NOTE: All runs and attempts are stored in memory.
        # This may be a memory issue for large runs.
        self._runs: Dict[str, TrainRun] = {}
        # {run_id: {attempt_id: TrainRunAttempt}}
        self._run_attempts: Dict[str, Dict[str, TrainRunAttempt]] = defaultdict(dict)

    def create_train_run(
        self,
        id: str,
        name: str,
        job_id: str,
        controller_actor_id: str,
        controller_log_file_path: str,
        run_config: RunConfig,
        train_loop_config: Optional[Dict],
        scaling_config: ScalingConfig,
        backend_config: BackendConfig,
        datasets: Dict[str, ray.data.Dataset],
        dataset_config: DataConfig,
    ) -> None:
        run_config_schema = RunConfigSchema(
            name=run_config.name,
            failure_config=FailureConfigSchema(
                max_failures=run_config.failure_config.max_failures,
                controller_failure_limit=run_config.failure_config.controller_failure_limit,
            ),
            worker_runtime_env=run_config.worker_runtime_env,
            checkpoint_config=CheckpointConfigSchema(
                num_to_keep=run_config.checkpoint_config.num_to_keep,
                checkpoint_score_attribute=run_config.checkpoint_config.checkpoint_score_attribute,
                checkpoint_score_order=run_config.checkpoint_config.checkpoint_score_order,
            ),
            storage_path=run_config.storage_path,
            storage_filesystem=run_config.storage_filesystem,
        )

        scaling_config_schema = ScalingConfigSchema(
            num_workers=scaling_config.num_workers,
            use_gpu=scaling_config.use_gpu,
            resources_per_worker=scaling_config.resources_per_worker,
            placement_strategy=scaling_config.placement_strategy,
            accelerator_type=scaling_config.accelerator_type,
            use_tpu=scaling_config.use_tpu,
            topology=scaling_config.topology,
            bundle_label_selector=scaling_config.label_selector,
        )

        backend_config_schema = BackendConfigSchema(
            framework=backend_config.framework,
            config=backend_config.to_dict(),
        )

        run_settings = RunSettings(
            train_loop_config=train_loop_config,
            backend_config=backend_config_schema,
            scaling_config=scaling_config_schema,
            datasets=list(datasets.keys()),
            data_config=construct_data_config(dataset_config),
            run_config=run_config_schema,
        )

        run = TrainRun(
            id=id,
            name=name,
            job_id=job_id,
            status=RunStatus.INITIALIZING,
            status_detail=None,
            controller_actor_id=controller_actor_id,
            start_time_ns=current_time_ns(),
            end_time_ns=None,
            controller_log_file_path=controller_log_file_path,
            framework_versions={"ray": ray.__version__},
            run_settings=run_settings,
        )
        self._runs[run.id] = run
        # Block so the initial run state isn't lost if the controller exits
        # right after. Without this, the .remote() task could still be in the
        # caller's outbound queue when the controller dies, leaving the state
        # actor with no record of the run.
        self._create_or_update_train_run(run, block=True)

    def update_train_run_scheduling(
        self,
        run_id: str,
        resize_decision: Optional[ResizeDecision] = None,
    ) -> None:
        if resize_decision is not None:
            status_detail = _get_scheduling_status_detail(
                resize_decision.num_workers, resize_decision.resources_per_worker
            )
        else:
            status_detail = None

        run = self._runs[run_id]
        run.status = RunStatus.SCHEDULING
        run.status_detail = status_detail
        self._create_or_update_train_run(run)

    def update_train_run_running(
        self,
        run_id: str,
    ) -> None:
        run = self._runs[run_id]
        run.status = RunStatus.RUNNING
        run.status_detail = None
        self._create_or_update_train_run(run)

    def update_train_run_restarting(
        self,
        run_id: str,
    ) -> None:
        run = self._runs[run_id]
        run.status = RunStatus.RESTARTING
        run.status_detail = None
        self._create_or_update_train_run(run)

    def update_train_run_resizing(
        self,
        run_id: str,
    ) -> None:
        run = self._runs[run_id]
        run.status = RunStatus.RESIZING
        run.status_detail = None
        self._create_or_update_train_run(run)

    def update_train_run_finished(
        self,
        run_id: str,
    ):
        run = self._runs[run_id]
        run.status = RunStatus.FINISHED
        run.status_detail = None
        run.end_time_ns = current_time_ns()
        # Block on terminal status so the final state isn't lost if the controller exits right after.
        self._create_or_update_train_run(run, block=True)

    def update_train_run_errored(
        self,
        run_id: str,
        status_detail: str,
    ):
        run = self._runs[run_id]
        run.status = RunStatus.ERRORED
        run.status_detail = status_detail
        run.end_time_ns = current_time_ns()
        # Block on terminal status so the final state isn't lost if the controller exits right after.
        self._create_or_update_train_run(run, block=True)

    def update_train_run_aborted(
        self,
        run_id: str,
    ):
        run = self._runs[run_id]
        update_train_run_aborted(run=run, graceful=True)
        # Block on terminal status so the final state isn't lost if the controller exits right after.
        self._create_or_update_train_run(run, block=True)

    def update_train_run_framework_versions(
        self, run_id: str, framework_versions: Dict[str, str]
    ):
        run = self._runs[run_id]
        run.framework_versions = framework_versions
        self._create_or_update_train_run(run)

    def create_train_run_attempt(
        self,
        run_id: str,
        attempt_id: str,
        num_workers: int,
        resources_per_worker: Dict[str, float],
    ) -> None:
        status_detail = _get_scheduling_status_detail(num_workers, resources_per_worker)
        resources = [
            TrainResources(resources=resources_per_worker) for _ in range(num_workers)
        ]
        run_attempt = TrainRunAttempt(
            run_id=run_id,
            attempt_id=attempt_id,
            start_time_ns=current_time_ns(),
            status=RunAttemptStatus.PENDING,
            status_detail=status_detail,
            resources=resources,
            workers=[],  # Not started yet.
        )

        self._run_attempts[run_id][attempt_id] = run_attempt
        self._create_or_update_train_run_attempt(run_attempt)

    def update_train_run_attempt_running(
        self, run_id: str, attempt_id: str, workers: List[Worker]
    ) -> None:
        def _convert_worker(worker: Worker) -> TrainWorker:

            actor: ActorHandle = worker.actor
            distributed_context: DistributedContext = worker.distributed_context
            actor_metadata: ActorMetadata = worker.metadata

            return TrainWorker(
                world_rank=distributed_context.world_rank,
                local_rank=distributed_context.local_rank,
                node_rank=distributed_context.node_rank,
                actor_id=actor._actor_id.hex(),
                node_id=actor_metadata.node_id,
                node_ip=actor_metadata.node_ip,
                pid=actor_metadata.pid,
                gpu_ids=actor_metadata.gpu_ids,
                status=ActorStatus.ALIVE,
                resources=TrainResources(resources=worker.resources),
                log_file_path=worker.log_file_path,
            )

        workers: List[TrainWorker] = [_convert_worker(worker) for worker in workers]

        run_attempt = self._run_attempts[run_id][attempt_id]
        run_attempt.status = RunAttemptStatus.RUNNING
        run_attempt.status_detail = None
        run_attempt.workers = workers
        self._create_or_update_train_run_attempt(run_attempt)

    def update_train_run_attempt_finished(
        self,
        run_id: str,
        attempt_id: str,
    ):
        run_attempt = self._run_attempts[run_id][attempt_id]
        run_attempt.status = RunAttemptStatus.FINISHED
        run_attempt.status_detail = None
        run_attempt.end_time_ns = current_time_ns()
        mark_workers_dead(run_attempt)
        # Block to avoid case where controller is dead but attempt is not terminal.
        self._create_or_update_train_run_attempt(run_attempt, block=True)

    def update_train_run_attempt_errored(
        self,
        run_id: str,
        attempt_id: str,
        status_detail: str,
    ):
        run_attempt = self._run_attempts[run_id][attempt_id]
        run_attempt.status = RunAttemptStatus.ERRORED
        run_attempt.status_detail = status_detail
        run_attempt.end_time_ns = current_time_ns()
        mark_workers_dead(run_attempt)
        # Block to avoid case where controller is dead but attempt is not terminal.
        self._create_or_update_train_run_attempt(run_attempt, block=True)

    def update_train_run_attempt_aborted(
        self,
        run_id: str,
        attempt_id: str,
    ):
        run_attempt = self._run_attempts[run_id][attempt_id]
        update_train_run_attempt_aborted(run_attempt=run_attempt, graceful=True)
        # Block to avoid case where controller is dead but attempt is not terminal.
        self._create_or_update_train_run_attempt(run_attempt, block=True)

    def get_train_run_framework(self, run_id: str) -> Optional[TrainingFramework]:
        run = self._runs[run_id]
        return run.run_settings.backend_config.framework

    def _create_or_update_train_run(
        self, run: TrainRun, *, block: bool = False
    ) -> None:
        ref = self._state_actor.create_or_update_train_run.remote(run)
        if block:
            ray.get(ref)

    def _create_or_update_train_run_attempt(
        self, run_attempt: TrainRunAttempt, *, block: bool = False
    ) -> None:
        ref = self._state_actor.create_or_update_train_run_attempt.remote(run_attempt)
        if block:
            ray.get(ref)


def _get_scheduling_status_detail(
    num_workers: int, resources_per_worker: Dict[str, float]
) -> str:
    return f"Scheduling {num_workers} workers, each requiring: {resources_per_worker}."