import math
from collections.abc import Mapping
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

from pydantic import field_validator

from ray._common.pydantic_compat import BaseModel, Field
from ray.dashboard.modules.job.pydantic_models import JobDetails
from ray.train.v2._internal.util import TrainingFramework
from ray.util.annotations import DeveloperAPI

MAX_ERROR_STACK_TRACE_LENGTH = 50000


def _to_json_serializable_value(value: Any, *, max_depth: int = 3) -> Any:
    """Recursively coerce a value into a human-readable, JSON serializable representation.

    If ``value`` is a list or dict, this function walks through it and replaces non-JSON
    serializable fields (e.g. custom objects, modules, tensors, callables, etc.) with a
    human-readable string representation.

    Args:
        value: Any Python value. Primitives pass through; collections recurse;
            other types are stringified.
        max_depth: Truncates dicts nested beyond ``max_depth`` to ``"..."``.
            Lists do not consume depth.

    Returns:
        The JSON serializable representation of the value.
    """
    if max_depth <= 0:
        raise ValueError("max_depth must be greater than 0")

    def _safe_str(v):
        try:
            return str(v)
        except Exception:
            return type(v).__name__

    def _walk(value, depth):
        if value is None or isinstance(value, (bool, int, str)):
            return value
        if isinstance(value, float):
            return str(value) if not math.isfinite(value) else value
        if isinstance(value, Mapping):
            if depth <= 0:
                return "..."
            try:
                items = list(value.items())
            except Exception:
                # Custom Mapping subclass with a broken `.items()`.
                return type(value).__name__
            return {_safe_str(k): _walk(v, depth - 1) for k, v in items}

        # Tuples, sets, and frozensets all become lists in JSON.
        if isinstance(value, (list, tuple, set, frozenset)):
            return [_walk(v, depth) for v in value]

        cls = type(value)
        # Use class name if no custom string representation is defined.
        if cls.__str__ is object.__str__ and cls.__repr__ is object.__repr__:
            return cls.__name__
        return _safe_str(value)

    return _walk(value, max_depth)


@DeveloperAPI
class RunStatus(str, Enum):
    """Enumeration of the possible statuses for a Train run."""

    # ====== Active States ======
    # The Train run is currently in the process of initializing.
    INITIALIZING = "INITIALIZING"
    # The Train run is waiting to be scheduled.
    SCHEDULING = "SCHEDULING"
    # The Train run is currently in progress.
    RUNNING = "RUNNING"
    # The Train run is recovering from a failure or restart.
    RESTARTING = "RESTARTING"
    # The Train run is resizing.
    RESIZING = "RESIZING"

    # ===== Terminal States ======
    # The Train run completed successfully.
    FINISHED = "FINISHED"
    # The Train run failed due to an error in the training workers.
    ERRORED = "ERRORED"
    # The Train run was terminated due to system or controller errors.
    ABORTED = "ABORTED"

    def is_terminal(self) -> bool:
        return self in [RunStatus.FINISHED, RunStatus.ERRORED, RunStatus.ABORTED]


@DeveloperAPI
class RunAttemptStatus(str, Enum):
    """Enumeration of the possible statuses for a Train run attempt."""

    # ====== Active States ======
    # The run attempt is waiting to be scheduled.
    PENDING = "PENDING"
    # The run attempt is currently in progress.
    RUNNING = "RUNNING"

    # ===== Terminal States =====
    # The run attempt completed successfully.
    FINISHED = "FINISHED"
    # The run attempt failed due to an error in the training workers.
    ERRORED = "ERRORED"
    # The run attempt was terminated due to system or controller errors.
    ABORTED = "ABORTED"

    def is_terminal(self) -> bool:
        return self in [
            RunAttemptStatus.FINISHED,
            RunAttemptStatus.ERRORED,
            RunAttemptStatus.ABORTED,
        ]


@DeveloperAPI
class ActorStatus(str, Enum):
    """Enumeration of the statuses for a Train worker actor."""

    # The actor is currently active.
    ALIVE = "ALIVE"
    # The actor is no longer active.
    DEAD = "DEAD"


@DeveloperAPI
class TrainResources(BaseModel):
    """Resources allocated for a Train worker or run."""

    resources: Dict[str, float] = Field(
        description="A dictionary specifying the types and amounts of resources "
        "allocated (e.g., CPU, GPU)."
    )


@DeveloperAPI
class TrainWorker(BaseModel):
    """Metadata about a Ray Train worker."""

    world_rank: int = Field(
        description="The global rank of the worker in the training cluster."
    )
    local_rank: int = Field(description="The local rank of the worker on its node.")
    node_rank: int = Field(description="The rank of the worker's node in the cluster.")
    actor_id: str = Field(description="The unique ID of the worker's actor.")
    node_id: str = Field(
        description="The unique ID of the node where the worker is running."
    )
    node_ip: str = Field(
        description="The IP address of the node where the worker is running."
    )
    pid: int = Field(description="The process ID of the worker.")
    gpu_ids: List[int] = Field(description="A list of GPU IDs allocated to the worker.")
    status: Optional[ActorStatus] = Field(
        None, description="The current status of the worker actor."
    )
    resources: TrainResources = Field(
        description="The resources allocated to this Train worker."
    )
    log_file_path: Optional[str] = Field(
        None, description="The path to the log file for the Train worker."
    )


@DeveloperAPI
class MemoryInfo(BaseModel):
    """Memory usage information for a process."""

    rss: int = Field(description="The resident set size (RSS) memory usage in bytes.")
    vms: int = Field(description="The virtual memory size (VMS) usage in bytes.")
    pfaults: Optional[int] = Field(None, description="The number of page faults.")
    pageins: Optional[int] = Field(None, description="The number of page-ins.")


@DeveloperAPI
class ProcessStats(BaseModel):
    """CPU and memory statistics for a process."""

    cpuPercent: float = Field(description="The percentage of CPU usage.")
    mem: Optional[List[int]] = Field(
        None,
        description="Memory statistics, including total memory, free memory, "
        "and memory usage ratio.",
    )
    memoryInfo: MemoryInfo = Field(description="Detailed memory usage information.")


class ProcessGPUUsage(BaseModel):
    """GPU usage statistics for a process."""

    pid: int = Field(description="The process ID.")
    gpuMemoryUsage: int = Field(description="The GPU memory usage in bytes.")


@DeveloperAPI
class GPUStats(BaseModel):
    """Statistics for a GPU."""

    uuid: str = Field(description="The unique identifier of the GPU.")
    index: int = Field(description="The index of the GPU.")
    name: str = Field(description="The name of the GPU.")
    utilizationGpu: Optional[float] = Field(
        None, description="The percentage utilization of the GPU."
    )
    memoryUsed: float = Field(description="The amount of GPU memory used in bytes.")
    memoryTotal: float = Field(description="The total amount of GPU memory in bytes.")
    processInfo: ProcessGPUUsage = Field(
        description="GPU usage statistics for the associated process."
    )


@DeveloperAPI
class DecoratedTrainWorker(TrainWorker):
    """Detailed metadata for a Ray Train worker, including process and GPU stats."""

    processStats: Optional[ProcessStats] = Field(
        None, description="CPU and memory statistics for the worker process."
    )
    gpus: List[GPUStats] = Field(
        default_factory=list,
        description="A list of GPUs used by the worker process,"
        " with detailed statistics.",
    )


@DeveloperAPI
class TrainRunAttempt(BaseModel):
    """Metadata for an individual attempt to execute a Train run."""

    run_id: str = Field(description="Unique identifier for the parent Train run.")
    attempt_id: str = Field(
        description="Unique identifier for this specific Train run attempt."
    )
    status: RunAttemptStatus = Field(
        description="The current execution status of the Train run attempt."
    )
    status_detail: Optional[str] = Field(
        None,
        description="Additional details about the status,"
        " including error messages if applicable.",
    )
    start_time_ns: int = Field(
        description="The UNIX timestamp (in nanoseconds)"
        " when the Train run attempt started."
    )
    end_time_ns: Optional[int] = Field(
        None,
        description="The UNIX timestamp (in nanoseconds)"
        " when the Train run attempt ended. "
        "If null, the attempt is still ongoing.",
    )
    resources: List[TrainResources] = Field(
        description="The resources (e.g., CPU, GPU) allocated to the Train run attempt."
    )
    workers: List[TrainWorker] = Field(
        description="List of Train workers participating in this attempt, "
        "sorted by global ranks."
    )


@DeveloperAPI
class DecoratedTrainRunAttempt(TrainRunAttempt):
    """Detailed metadata for a Train run attempt, including decorated worker data."""

    workers: List[DecoratedTrainWorker] = Field(
        description="A list of Train workers with detailed statistics, "
        "sorted by global ranks."
    )


@DeveloperAPI
class ExecutionOptions(BaseModel):
    """ExecutionOptions for a single Ray Data ingest pipeline."""

    resource_limits: Dict[str, Any] = Field(
        description="The resource limits applied to the Ray Data execution plan."
    )
    exclude_resources: Dict[str, Any] = Field(
        description="The resources excluded from the Ray Data execution plan "
        "(e.g. resources reserved by Ray Train workers)."
    )

    @field_validator("resource_limits", "exclude_resources", mode="before")
    @classmethod
    def _sanitize_dict(cls, v):
        return _to_json_serializable_value(v)

    preserve_order: bool = Field(
        description="Whether to preserve the order of outputs across operators."
    )
    actor_locality_enabled: bool = Field(
        description="Whether actor-based locality optimizations are enabled."
    )
    verbose_progress: bool = Field(
        description="Whether verbose progress reporting is enabled."
    )


@DeveloperAPI
class DataExecutionOptions(BaseModel):
    """ExecutionOptions for a Ray Train run, split into defaults and per-dataset overrides."""

    default: ExecutionOptions = Field(
        description="Execution options applied to any dataset without a per-dataset override."
    )
    per_dataset_execution_options: Dict[str, ExecutionOptions] = Field(
        default_factory=dict,
        description="Per-dataset execution option overrides, keyed by dataset name.",
    )


@DeveloperAPI
class DataConfig(BaseModel):
    """Configuration for dataset splitting and execution options within Ray Train."""

    datasets_to_split: Union[Literal["all"], List[str]] = Field(
        description="Which datasets to split; either 'all' or a list of dataset names."
    )
    execution_options: Optional[Dict] = Field(
        default=None,
        deprecated="DEPRECATED: Use data_execution_options instead.",
    )
    data_execution_options: DataExecutionOptions = Field(
        description="Data execution options"
    )
    enable_shard_locality: bool = Field(
        description="Whether to enable shard locality optimization."
    )


@DeveloperAPI
class ScalingConfig(BaseModel):
    """Scaling config for a Train run."""

    num_workers: Union[int, Tuple[int, int]] = Field(
        description="The number of workers for the Train run."
    )
    use_gpu: bool = Field(description="Whether to use GPUs for the Train run.")
    resources_per_worker: Optional[Dict[str, float]] = Field(
        None, description="The resources per worker for a Train run."
    )
    placement_strategy: str = Field(
        description="The placement strategy for the Train run."
    )
    accelerator_type: Optional[str] = Field(
        None, description="The accelerator type for the Train run."
    )
    use_tpu: bool = Field(description="Whether to use TPUs for the Train run.")
    topology: Optional[str] = Field(None, description="The topology for the Train run.")
    bundle_label_selector: Optional[
        Union[Dict[str, str], List[Dict[str, str]]]
    ] = Field(None, description="The bundle label selector for the Train run.")


@DeveloperAPI
class FailureConfig(BaseModel):
    """Failure config for a Train run."""

    max_failures: int = Field(
        description="The maximum number of failures for a Train run."
    )
    controller_failure_limit: int = Field(
        description="The maximum number of controller failures to tolerate."
    )


@DeveloperAPI
class CheckpointConfig(BaseModel):
    """Checkpoint config for a Train run."""

    num_to_keep: Optional[int] = Field(
        None,
        description="The number of most recent checkpoints to keep. Older checkpoints may be deleted.",
    )
    checkpoint_score_attribute: Optional[str] = Field(
        None,
        description="Attribute used to score and rank checkpoints; can be a metric key or attribute.",
    )
    checkpoint_score_order: Literal["max", "min"] = Field(
        description="Order to rank checkpoint scores, 'max' for higher-is-better, 'min' for lower-is-better.",
    )


@DeveloperAPI
class RunConfig(BaseModel):
    """Run configuration parameters for a Train run, encompassing failure,
    runtime environment, checkpoint settings, and storage path."""

    name: str = Field(description="The name of the Train run.")
    failure_config: FailureConfig = Field(
        description="The failure config for a Train run."
    )
    worker_runtime_env: Dict[str, Any] = Field(
        description="The worker runtime env for a Train run."
    )

    @field_validator("worker_runtime_env", mode="before")
    @classmethod
    def _sanitize_worker_runtime_env(cls, v):
        return _to_json_serializable_value(v)

    checkpoint_config: CheckpointConfig = Field(
        description="The checkpoint config for a Train run."
    )
    storage_path: str = Field(description="The storage path for a Train run.")
    storage_filesystem: Optional[str] = Field(
        None, description="The storage filesystem for a Train run."
    )

    @field_validator("storage_filesystem", mode="before")
    @classmethod
    def _sanitize_storage_filesystem(cls, v):
        return _to_json_serializable_value(v)


@DeveloperAPI
class BackendConfig(BaseModel):
    """Backend config for a Train run."""

    framework: Optional[TrainingFramework] = Field(
        None, description="The training framework for this backend config."
    )
    config: Dict[str, Any] = Field(
        description="Training framework-specific configuration fields."
    )

    @field_validator("config", mode="before")
    @classmethod
    def _sanitize_config(cls, v):
        return _to_json_serializable_value(v)


@DeveloperAPI
class RunSettings(BaseModel):
    """Settings for a Train run, primarily consisting of configs set before a train run starts.

    This includes the train loop config, backend config, scaling config, dataset configs,
    and runtime configuration.
    """

    train_loop_config: Optional[Dict] = Field(
        None, description="The user defined train loop config for a Train run."
    )

    @field_validator("train_loop_config", mode="before")
    @classmethod
    def _sanitize_train_loop_config(cls, v):
        return _to_json_serializable_value(v)

    backend_config: BackendConfig = Field(
        description="The backend config for a Train run. Can vary with the framework (e.g. TorchConfig)"
    )
    scaling_config: ScalingConfig = Field(
        description="The scaling config for this Train run."
    )
    datasets: List[str] = Field(
        description="A list of dataset names for a Train run.",
    )
    data_config: DataConfig = Field(
        description="The data config for a Train run.",
    )
    run_config: RunConfig = Field(
        description="Run configuration for this Train run, including failure, runtime environment, checkpoint settings, and storage path."
    )


@DeveloperAPI
class TrainRun(BaseModel):
    """Metadata for a Ray Train run, including its details and status."""

    id: str = Field(description="Unique identifier for the Train run.")
    name: str = Field(description="Human-readable name assigned to the Train run.")
    job_id: str = Field(description="The Ray Job ID associated with this Train run.")
    controller_actor_id: str = Field(
        description="Unique ID of the actor managing the Train run."
    )
    status: RunStatus = Field(
        description="The current execution status of the Train run."
    )
    status_detail: Optional[str] = Field(
        None,
        description="Additional details about the current status, "
        "including error messages if applicable.",
    )
    start_time_ns: int = Field(
        description="The UNIX timestamp (in nanoseconds) when the Train run started."
    )
    end_time_ns: Optional[int] = Field(
        None,
        description="The UNIX timestamp (in nanoseconds) when the Train run ended. "
        "If null, the run is still in progress.",
    )
    controller_log_file_path: Optional[str] = Field(
        None, description="The path to the log file for the Train run controller."
    )
    framework_versions: Dict[str, str] = Field(
        description="The relevant framework versions for this Train run,"
        "including the Ray version and training framework version."
    )
    run_settings: RunSettings = Field(
        description="The run settings for this Train run, including train loop config, "
        "backend config, scaling config, dataset details, and runtime configuration."
    )


@DeveloperAPI
class DecoratedTrainRun(TrainRun):
    """Detailed metadata for a Ray Train run, including attempts and job details."""

    attempts: List[DecoratedTrainRunAttempt] = Field(
        description="A list of attempts made to execute the Train run."
    )
    job_details: Optional[JobDetails] = Field(
        None,
        description="Detailed information about the job that initiated this Train run.",
    )


@DeveloperAPI
class TrainRunsResponse(BaseModel):
    """Response containing a list of decorated Train runs."""

    train_runs: List[DecoratedTrainRun] = Field(
        description="A list of Train runs with detailed metadata."
    )