import time
from typing import TYPE_CHECKING, Dict, List, Optional, Union

from lightning_sdk.api.job_api import JobApiV2, V1ClusterAccelerator, _keep_minutes_spec_kwargs
from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
from lightning_sdk.api.utils import (
    _machine_to_compute_name,
    cached_lightning_client,
    resolve_path_mappings,
    resolve_tags,
)
from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
from lightning_sdk.lightning_cloud.openapi import (
    JobsServiceCreateMultiMachineJobBody,
    JobsServiceUpdateMultiMachineJobBody,
    V1EnvVar,
    V1Job,
    V1JobSpec,
    V1MultiMachineJob,
    V1MultiMachineJobFaultTolerance,
    V1MultiMachineJobFaultToleranceStrategy,
    V1MultiMachineJobState,
)
from lightning_sdk.machine import Machine

if TYPE_CHECKING:
    from lightning_sdk.status import Status


class MMTApiV2:
    """Native (v2 jobs-service) API client for multi-machine training jobs."""

    def __init__(self) -> None:
        self._cloud_url = _cloud_url()
        self._client = cached_lightning_client()

    def submit_job(
        self,
        name: str,
        num_machines: int,
        command: Optional[str],
        cloud_account: Optional[str],
        teamspace_id: str,
        studio_id: Optional[str],
        image: Optional[str],
        machine: Union[Machine, str],
        interruptible: bool,
        env: Optional[Dict[str, str]],
        image_credentials: Optional[str],
        cloud_account_auth: bool,
        entrypoint: Optional[str],
        path_mappings: Optional[Dict[str, str]],
        max_runtime: Optional[int],
        reuse_snapshot: bool,
        placement_group_id: Optional[str] = None,
        scratch_disks: Optional[Dict[str, int]] = None,
        max_run_attempts: Optional[int] = None,
        tags: Optional[List[str]] = None,
        keep_minutes: Optional[int] = None,
    ) -> V1MultiMachineJob:
        """Submit a v2 multi-machine job and return the created job object.

        Args:
            name: The name to assign to the new multi-machine job.
            num_machines: The number of machines to allocate for the job.
            command: The shell command each machine will execute, or ``None`` when using a custom entrypoint.
            cloud_account: The cloud account identifier to use for compute, or ``None`` to use the default.
            teamspace_id: The ID of the teamspace that will own the job.
            studio_id: The ID of the Studio from which the job is launched, or ``None``.
            image: The container image to run, or ``None`` to use the Studio snapshot.
            machine: The machine type (as a ``Machine`` enum or a string slug) to run the job on.
            interruptible: Whether the job may run on interruptible (spot) instances.
            env: Optional mapping of environment variable names to values.
            image_credentials: Name of the secret holding image-pull credentials, or ``None``.
            cloud_account_auth: Whether to pass cloud-account credentials into the container.
            entrypoint: The entrypoint command used to launch the job process.
            path_mappings: Optional mapping of local paths to remote artifact destinations.
            max_runtime: DWS (Dynamic Workload Scheduler) reservation duration in seconds
                (e.g. some top-end GCP GPUs or Lightning baremetal). Not a time limit; see ``Job.run``.
                ``None`` means no reservation is requested.
            reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot.
            placement_group_id: Optional placement group identifier for colocating the job.
            scratch_disks: Not supported for multi-machine jobs. Kept for parity with ``JobApiV2.submit_job``.
            max_run_attempts: Max number of run attempts for this multi-machine job. ``None`` or ``0``
                means unset (legacy, no retries); ``1`` means a single attempt (no retries). When
                greater than ``1`` the ``RECREATE_ALL_NODES`` fault tolerance strategy is set
                automatically on the multi-machine job body. Set at the multi-machine job level,
                not on the per-machine ``JobSpec``.
            tags: Optional teamspace tag names to apply to the job. Tags that don't exist in the
                teamspace yet are created, which requires permission to create tags.
            keep_minutes: Minutes to keep the machines after the job stops. ``None`` or ``0``
                releases the machines immediately.

        Returns:
            The newly created ``V1MultiMachineJob`` object.
        """
        if scratch_disks:
            raise ValueError("scratch_disks are not supported for multi-machine jobs")
        body = self._create_mmt_body(
            name=name,
            num_machines=num_machines,
            command=command,
            cloud_account=cloud_account,
            studio_id=studio_id,
            image=image,
            machine=machine,
            interruptible=interruptible,
            env=env,
            image_credentials=image_credentials,
            cloud_account_auth=cloud_account_auth,
            entrypoint=entrypoint,
            path_mappings=path_mappings,
            max_runtime=max_runtime,
            reuse_snapshot=reuse_snapshot,
            placement_group_id=placement_group_id,
            max_run_attempts=max_run_attempts,
            tags=tags,
            keep_minutes=keep_minutes,
        )

        job: V1MultiMachineJob = self._client.jobs_service_create_multi_machine_job(project_id=teamspace_id, body=body)
        return job

    @staticmethod
    def _create_mmt_body(
        name: str,
        num_machines: int,
        command: Optional[str],
        cloud_account: Optional[str],
        studio_id: Optional[str],
        image: Optional[str],
        machine: Union[Machine, str],
        interruptible: bool,
        env: Optional[Dict[str, str]],
        image_credentials: Optional[str],
        cloud_account_auth: bool,
        entrypoint: Optional[str],
        path_mappings: Optional[Dict[str, str]],
        reuse_snapshot: bool,
        max_runtime: Optional[int] = None,
        machine_image_version: Optional[str] = None,
        placement_group_id: Optional[str] = None,
        max_run_attempts: Optional[int] = None,
        tags: Optional[List[str]] = None,
        keep_minutes: Optional[int] = None,
    ) -> JobsServiceCreateMultiMachineJobBody:
        """Build the request body for creating a v2 multi-machine job.

        Args:
            name: The name to assign to the new multi-machine job.
            num_machines: The number of machines to allocate for the job.
            command: The shell command each machine will execute, or ``None`` when using a custom entrypoint.
            cloud_account: The cloud account identifier to use for compute, or ``None`` to use the default.
            studio_id: The ID of the Studio from which the job is launched, or ``None``.
            image: The container image to run, or ``None`` to use the Studio snapshot.
            machine: The machine type (as a ``Machine`` enum or a string slug) to run the job on.
            interruptible: Whether the job may run on interruptible (spot) instances.
            env: Optional mapping of environment variable names to values.
            image_credentials: Name of the secret holding image-pull credentials, or ``None``.
            cloud_account_auth: Whether to pass cloud-account credentials into the container.
            entrypoint: The entrypoint command used to launch the job process.
            path_mappings: Optional mapping of local paths to remote artifact destinations.
            reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot.
            max_runtime: DWS (Dynamic Workload Scheduler) reservation duration in seconds
                (e.g. some top-end GCP GPUs or Lightning baremetal). Not a time limit; see ``Job.run``.
                ``None`` means no reservation is requested.
            machine_image_version: Pinned machine-image version string, or ``None`` for the default.
            placement_group_id: Optional placement group identifier for colocating the job.
            max_run_attempts: Max number of run attempts for this multi-machine job. ``None`` or ``0``
                means unset (legacy, no retries); ``1`` means a single attempt (no retries). When
                greater than ``1`` the ``RECREATE_ALL_NODES`` fault tolerance strategy is set
                automatically on the multi-machine job body. Set at the multi-machine job level,
                not on the per-machine ``JobSpec``.
            tags: Optional teamspace tag names to apply to the job.
            keep_minutes: Minutes to keep the machines after the job stops. ``None`` or ``0``
                releases the machines immediately.

        Returns:
            A fully populated ``JobsServiceCreateMultiMachineJobBody`` ready to be sent to the jobs service.
        """
        env_vars = []
        if env is not None:
            for k, v in env.items():
                env_vars.append(V1EnvVar(name=k, value=v))

        instance_name = _machine_to_compute_name(machine)

        run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if (studio_id is not None and reuse_snapshot) else ""

        path_mappings_list = resolve_path_mappings(mappings=path_mappings or {})

        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
        optional_spec_kwargs = {}
        if max_runtime:
            optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
        optional_spec_kwargs.update(_keep_minutes_spec_kwargs(keep_minutes))

        spec = V1JobSpec(
            cloudspace_id=studio_id or "",
            cluster_id=cloud_account or "",
            command=command or "",
            entrypoint=entrypoint,
            env=env_vars,
            image=image or "",
            instance_name=instance_name,
            run_id=run_id,
            spot=interruptible,
            image_cluster_credentials=cloud_account_auth,
            image_secret_ref=image_credentials or "",
            path_mappings=path_mappings_list,
            machine_image_version=machine_image_version,
            placement_group_id=placement_group_id,
            **optional_spec_kwargs,
        )
        # max_run_attempts and fault_tolerance live on the multi-machine job body,
        # NOT on the per-machine JobSpec. Fault tolerance is auto-derived: when the
        # job is multi-machine and allows retries (max_run_attempts > 1) the only
        # currently supported strategy (RECREATE_ALL_NODES) is set automatically;
        # otherwise the field is left unset.
        fault_tolerance = None
        if num_machines > 1 and (max_run_attempts or 0) > 1:
            fault_tolerance = V1MultiMachineJobFaultTolerance(
                strategy=V1MultiMachineJobFaultToleranceStrategy.RECREATE_ALL_NODES
            )
        return JobsServiceCreateMultiMachineJobBody(
            name=name,
            spec=spec,
            cluster_id=cloud_account or "",
            machines=num_machines,
            max_run_attempts=max_run_attempts,
            fault_tolerance=fault_tolerance,
            tags=resolve_tags(tags),
        )

    def get_job_by_name(self, name: str, teamspace_id: str) -> V1MultiMachineJob:
        """Fetch a multi-machine job by its unique name within a teamspace.

        Args:
            name: The name of the multi-machine job to look up.
            teamspace_id: The ID of the teamspace that owns the job.

        Returns:
            The matching ``V1MultiMachineJob`` object.
        """
        job: V1MultiMachineJob = self._client.jobs_service_get_multi_machine_job_by_name(
            project_id=teamspace_id, name=name
        )
        return job

    def get_job(self, job_id: str, teamspace_id: str) -> V1MultiMachineJob:
        """Fetch a multi-machine job by its unique ID.

        Args:
            job_id: The unique identifier of the multi-machine job to retrieve.
            teamspace_id: The ID of the teamspace that owns the job.

        Returns:
            The matching ``V1MultiMachineJob`` object.
        """
        job: V1MultiMachineJob = self._client.jobs_service_get_multi_machine_job(project_id=teamspace_id, id=job_id)
        return job

    def stop_job(self, job_id: str, teamspace_id: str) -> None:
        """Request to stop a running multi-machine job and wait until it reaches a terminal state.

        Args:
            job_id: The unique identifier of the multi-machine job to stop.
            teamspace_id: The ID of the teamspace that owns the job.
        """
        from lightning_sdk.status import Status

        current_job = self.get_job(job_id=job_id, teamspace_id=teamspace_id)

        current_state = self._job_state_to_external(current_job.desired_state)

        if current_state in (
            Status.Stopped,
            Status.Completed,
            Status.Failed,
        ):
            return

        if current_state != Status.Stopped:
            update_body = JobsServiceUpdateMultiMachineJobBody(desired_state=V1MultiMachineJobState.STOP)
            self._client.jobs_service_update_multi_machine_job(body=update_body, project_id=teamspace_id, id=job_id)

        while True:
            current_job = self.get_job(job_id=job_id, teamspace_id=teamspace_id)
            if self._job_state_to_external(current_job.state) in (
                Status.Stopping,
                Status.Completed,
                Status.Stopped,
                Status.Failed,
            ):
                break
            time.sleep(1)

    def set_tags(self, job_id: str, teamspace_id: str, tags: List[str]) -> None:
        """Replace the multi-machine job's tags with the given tag names.

        Tags that don't exist in the teamspace yet are created.
        Pass an empty list to remove all tags.

        Args:
            job_id: The unique identifier of the multi-machine job.
            teamspace_id: The ID of the teamspace that owns the job.
            tags: The list of tag names to set on the job.
        """
        from lightning_sdk.lightning_cloud.openapi import JobsServiceSetWorkloadTagsBody

        body = JobsServiceSetWorkloadTagsBody(tags=tags)
        self._client.jobs_service_set_workload_tags(
            body=body,
            project_id=teamspace_id,
            workload_type="multi_machine_job",
            workload_id=job_id,
        )

    def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str] = None) -> None:
        """Permanently delete a multi-machine job.

        Args:
            job_id: The unique identifier of the multi-machine job to delete.
            teamspace_id: The ID of the teamspace that owns the job.
            cloudspace_id: Ignored. Kept for parity with ``JobApiV2.delete_job``.
        """
        self._client.jobs_service_delete_multi_machine_job(project_id=teamspace_id, id=job_id)

    def list_mmt_subjobs(self, job_id: str, teamspace_id: str) -> List[V1Job]:
        """Return the individual machine sub-jobs that make up a multi-machine job.

        Args:
            job_id: The unique identifier of the parent multi-machine job.
            teamspace_id: The ID of the teamspace that owns the job.

        Returns:
            A list of ``V1Job`` objects representing each per-machine sub-job.
        """
        jobs_resp = self._client.jobs_service_list_jobs(project_id=teamspace_id, multi_machine_job_id=job_id)
        return jobs_resp.jobs

    def _job_state_to_external(self, state: V1MultiMachineJobState) -> "Status":
        """Convert a raw ``V1MultiMachineJobState`` enum value to the public ``Status`` enum.

        Args:
            state: The raw ``V1MultiMachineJobState`` value returned by the jobs service.

        Returns:
            The corresponding public ``Status`` enum value, defaulting to ``Status.Pending``
            for any unrecognised state.
        """
        from lightning_sdk.status import Status

        if str(state) == V1MultiMachineJobState.UNSPECIFIED:
            return Status.Pending
        if str(state) == V1MultiMachineJobState.RUNNING:
            return Status.Running
        if str(state) == V1MultiMachineJobState.STOPPED:
            return Status.Stopped
        if str(state) == V1MultiMachineJobState.COMPLETED:
            return Status.Completed
        if str(state) == V1MultiMachineJobState.FAILED:
            return Status.Failed
        if str(state) in (V1MultiMachineJobState.STOP, V1MultiMachineJobState.DELETE):
            return Status.Stopping
        if str(state) == V1MultiMachineJobState.DELETED:
            return Status.Stopped
        return Status.Pending

    def get_studio_name(self, job: V1MultiMachineJob) -> Optional[str]:
        """Return the name of the Studio linked to this job, or ``None`` if none is attached.

        Args:
            job: The multi-machine job whose linked Studio name is resolved.

        Returns:
            The display name of the Studio, or ``None`` if the job has no associated Studio.
        """
        if not job.spec.cloudspace_id:
            return None
        return JobApiV2._cached_studio_name(job.project_id, job.spec.cloudspace_id)

    def get_image_name(self, job: V1MultiMachineJob) -> Optional[str]:
        """Return the container image used by this job, or ``None`` if not set.

        Args:
            job: The multi-machine job whose image name is retrieved.

        Returns:
            The container image string, or ``None`` if no image was specified.
        """
        return job.spec.image or None

    def get_command(self, job: V1MultiMachineJob) -> str:
        """Return the shell command that the job executes.

        Args:
            job: The multi-machine job whose command is retrieved.

        Returns:
            The shell command string stored in the job spec.
        """
        return job.spec.command

    def _get_job_machine_from_spec(self, spec: V1JobSpec, teamspace_id: str, org_id: str) -> "Machine":
        """Resolve the ``Machine`` object from a job spec by matching against available accelerators.

        Args:
            spec: The job spec containing the instance name and cluster identifier.
            teamspace_id: The ID of the teamspace used for the accelerator lookup.
            org_id: The organisation ID used for the accelerator lookup.

        Returns:
            The ``Machine`` enum value that matches the spec's instance, falling back to
            ``Machine.from_str`` if no accelerator record matches.
        """
        predefined = Machine._predefined_from_str(spec.instance_name or spec.instance_type)
        if predefined is not None:
            return predefined

        accelerators = self._get_machines_for_cloud_account(
            teamspace_id=teamspace_id,
            cloud_account_id=spec.cluster_id,
            org_id=org_id,
        )

        for accelerator in accelerators:
            possible_identifiers = (
                accelerator.slug,
                accelerator.slug_multi_cloud,
                accelerator.instance_id,
            )
            if (spec.instance_name and spec.instance_name in possible_identifiers) or (
                spec.instance_type and spec.instance_type in possible_identifiers
            ):
                return Machine._from_accelerator(accelerator)

        return Machine.from_str(spec.instance_name or spec.instance_type)

    def _get_machines_for_cloud_account(
        self, teamspace_id: str, cloud_account_id: str, org_id: str
    ) -> List[V1ClusterAccelerator]:
        """Return only the enabled accelerators for a given cloud account.

        Args:
            teamspace_id: The ID of the teamspace used for the accelerator lookup.
            cloud_account_id: The cloud account whose accelerators are queried.
            org_id: The organisation ID used for the accelerator lookup.

        Returns:
            A list of ``V1ClusterAccelerator`` objects that are marked as enabled.
        """
        return JobApiV2._get_machines_for_cloud_account(
            teamspace_id=teamspace_id,
            cloud_account_id=cloud_account_id,
            org_id=org_id,
        )

    def get_total_cost(self, job: V1MultiMachineJob) -> float:
        """Return the accumulated cost for this multi-machine job in USD.

        Args:
            job: The multi-machine job whose cost is retrieved.

        Returns:
            The total cost incurred by the job, expressed in US dollars.
        """
        return job.total_cost

    def get_num_machines(self, job: V1MultiMachineJob) -> int:
        """Return the number of machines allocated for this multi-machine job.

        Args:
            job: The multi-machine job whose machine count is retrieved.

        Returns:
            The number of machines allocated for the job.
        """
        return job.machines