import time
import warnings
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from urllib.request import urlopen

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,
    remove_datetime_prefix,
    resolve_path_mappings,
    resolve_tags,
)
from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
from lightning_sdk.lightning_cloud.openapi import (
    JobsServiceCreateJobBody,
    JobsServiceUpdateJobBody,
    V1CloudProvider,
    V1ClusterAccelerator,
    V1DownloadJobLogsResponse,
    V1EnvVar,
    V1Job,
    V1JobSpec,
    V1Volume,
)
from lightning_sdk.lightning_cloud.openapi.rest import ApiException
from lightning_sdk.machine import Machine

if TYPE_CHECKING:
    from lightning_sdk.status import Status


def _keep_minutes_spec_kwargs(keep_minutes: Optional[int]) -> Dict[str, Any]:
    """Map a user timeout onto JobSpec fields sent to the backend.

    ``None`` or ``0`` omits both fields (release the machine on stop). A value
    greater than 0 sets ``reserve_machines_timeout_minutes`` and, until the
    backend drops the compatibility bit, also ``keep_machine_after_stop=True``.
    """
    if not keep_minutes:
        return {}
    if keep_minutes < 0:
        raise ValueError("keep_minutes must be >= 0")
    return {
        "reserve_machines_timeout_minutes": keep_minutes,
        "keep_machine_after_stop": True,
    }


class JobApiV2:
    """Native (v2 jobs-service) API client for single-machine Jobs."""

    # these are stages the job can be in.
    v2_job_state_pending = "pending"
    v2_job_state_creating = "creating"
    v2_job_state_running = "running"
    v2_job_state_restarting = "restarting"
    v2_job_state_stopped = "stopped"
    v2_job_state_completed = "completed"
    v2_job_state_failed = "failed"
    v2_job_state_deleted = "deleted"
    v2_job_state_stopping = "stopping"

    # the backend holds these while the machine shuts down and the job's logs and artifacts
    # are still being collected.
    v2_job_state_complete = "complete"
    v2_job_state_fail = "fail"
    v2_job_state_delete = "delete"

    # this is the user action to stop the job.
    v2_job_state_stop = "stop"

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

    def submit_job(
        self,
        name: str,
        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] = None,
        max_run_attempts: Optional[int] = None,
        reuse_snapshot: bool = True,
        scratch_disks: Optional[Dict[str, int]] = None,
        placement_group_id: Optional[str] = None,
        num_machines: int = 1,
        tags: Optional[List[str]] = None,
        keep_minutes: Optional[int] = None,
    ) -> V1Job:
        """Submit a v2 job and return the created job object.

        Args:
            name: The name to assign to the new job.
            command: The shell command the job 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.
            max_run_attempts: Max number of run attempts for this job. ``None`` or ``0`` means
                unset (backend default). ``1`` means a single attempt (no retries).
            reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot.
            scratch_disks: Optional mapping of scratch-disk mount paths to their sizes in GiB.
            placement_group_id: Optional placement group identifier for colocating the job.
            num_machines: Must be 1 for single-machine jobs. Kept for parity with ``MMTApiV2.submit_job``.
            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 machine after the job stops. ``None`` or ``0``
                releases the machine immediately.

        Returns:
            The newly created ``V1Job`` object.
        """
        if num_machines != 1:
            raise ValueError("JobApiV2 only supports single-machine jobs (num_machines=1)")
        if scratch_disks is not None:
            sanitized_scratch_disks = {}
            for k, v in scratch_disks.items():
                sanitized_k = k if k.startswith("/teamspace/scratch/") else f"/teamspace/scratch/{k}"
                sanitized_scratch_disks[sanitized_k] = v
        else:
            sanitized_scratch_disks = None

        body = self._create_job_body(
            name=name,
            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,
            max_run_attempts=max_run_attempts,
            reuse_snapshot=reuse_snapshot,
            scratch_disks=sanitized_scratch_disks,
            placement_group_id=placement_group_id,
            tags=tags,
            keep_minutes=keep_minutes,
        )

        job: V1Job = self._client.jobs_service_create_job(project_id=teamspace_id, body=body)
        return job

    @staticmethod
    def _create_job_body(
        name: str,
        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,
        max_run_attempts: Optional[int] = None,
        machine_image_version: Optional[str] = None,
        scratch_disks: Optional[Dict[str, int]] = None,
        placement_group_id: Optional[str] = None,
        tags: Optional[List[str]] = None,
        keep_minutes: Optional[int] = None,
    ) -> JobsServiceCreateJobBody:
        """Build the request body for creating a v2 job.

        Args:
            name: The name to assign to the new job.
            command: The shell command the job 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.
            max_run_attempts: Max number of run attempts for this job. ``None`` or ``0`` means
                unset (backend default). ``1`` means a single attempt (no retries).
            machine_image_version: Pinned machine-image version string, or ``None`` for the default.
            scratch_disks: Optional mapping of scratch-disk mount paths to their sizes in GiB.
            placement_group_id: Optional placement group identifier for colocating the job.
            tags: Optional teamspace tag names to apply to the job.
            keep_minutes: Minutes to keep the machine after the job stops. ``None`` or ``0``
                releases the machine immediately.

        Returns:
            A fully populated ``JobsServiceCreateJobBody`` 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: Dict[str, Any] = {}
        if max_runtime:
            optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
        if max_run_attempts:
            optional_spec_kwargs["max_run_attempts"] = max_run_attempts
        optional_spec_kwargs.update(_keep_minutes_spec_kwargs(keep_minutes))

        # don't do default dicts, as they'll be mutable. Create a fresh one here
        scratch_disks = scratch_disks or {}

        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,
            volumes=[V1Volume(path=k, size_gb=v, ephemeral=True) for k, v in scratch_disks.items()],
            **optional_spec_kwargs,
        )
        return JobsServiceCreateJobBody(name=name, spec=spec, tags=resolve_tags(tags))

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

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

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

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

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

        Returns:
            The matching ``V1Job`` object.
        """
        job: V1Job = self._client.jobs_service_get_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 v2 job and wait until it reaches a terminal state.

        Args:
            job_id: The unique identifier of the 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.state)

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

        if current_state != Status.Stopping:
            update_body = JobsServiceUpdateJobBody(state=self.v2_job_state_stop)
            self._client.jobs_service_update_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.Stopped,
                Status.Completed,
                Status.Stopped,
                Status.Failed,
            ):
                break
            time.sleep(1)

    def rename_job(self, job_id: str, teamspace_id: str, new_name: str) -> V1Job:
        """Rename a standalone job (not part of a deployment, MMT, or pipeline).

        Args:
            job_id: The unique identifier of the job to rename.
            teamspace_id: The ID of the teamspace that owns the job.
            new_name: The new name for the job.

        Returns:
            The updated ``V1Job`` object.
        """
        update_body = JobsServiceUpdateJobBody(name=new_name)
        job: V1Job = self._client.jobs_service_update_job(body=update_body, project_id=teamspace_id, id=job_id)
        return job

    def set_tags(self, job_id: str, teamspace_id: str, tags: List[str]) -> None:
        """Replace the 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 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="job",
            workload_id=job_id,
        )

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

        Args:
            job_id: The unique identifier of the job to delete.
            teamspace_id: The ID of the teamspace that owns the job.
            cloudspace_id: The ID of the Studio (cloudspace) associated with the job, or ``None``.
        """
        self._client.jobs_service_delete_job(project_id=teamspace_id, id=job_id, cloudspace_id=cloudspace_id or "")

    def get_logs_finished(self, job_id: str, teamspace_id: str, timestamps: bool = False) -> str:
        """Download and return the saved log file of a finished job.

        Only used for jobs whose lines are not in the logs API (jobs that finished before
        logs moved to the current storage format).

        Args:
            job_id: The unique identifier of the job whose logs are fetched.
            teamspace_id: The ID of the teamspace that owns the job.
            timestamps: Keep each line's ``[<RFC3339>] `` prefix instead of stripping it.

        Returns:
            The decoded log text, with datetime prefixes stripped unless ``timestamps``.
        """
        resp: V1DownloadJobLogsResponse = self._client.jobs_service_download_job_logs(
            project_id=teamspace_id, id=job_id
        )

        data = urlopen(resp.url).read().decode("utf-8")
        if timestamps:
            return str(data)
        return remove_datetime_prefix(str(data))

    @staticmethod
    @lru_cache(maxsize=None)
    def _cached_studio_name(project_id: str, cloudspace_id: str) -> Optional[str]:
        """Resolve a studio display name, caching across Job/MMT API instances."""
        try:
            return (
                cached_lightning_client()
                .cloud_space_service_get_cloud_space(project_id=project_id, id=cloudspace_id)
                .name
            )
        except ApiException as ex:
            if ex.status == 404:
                return None
            raise

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

        Args:
            job: The v2 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 self._cached_studio_name(job.project_id, job.spec.cloudspace_id)

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

        Args:
            job: The v2 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: V1Job) -> str:
        """Return the shell command that the job executes.

        Args:
            job: The v2 job whose command is retrieved.

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

    def get_mmt_name(self, job: V1Job) -> str:
        """Return the parent multi-machine job name by stripping the machine-index suffix.

        Args:
            job: The v2 job to inspect.

        Returns:
            The parent multi-machine job name derived by removing the trailing index suffix,
            or an empty string if the job is not part of a multi-machine job.
        """
        if job.multi_machine_job_id:
            splits = job.name.rsplit("-", 1)
            if len(splits) == 2:
                return splits[0]
        return ""

    def _is_job_finished(self, job_id: str, teamspace_id: str) -> bool:
        """Return whether the job has reached a terminal state.

        Used while following logs to decide when to stop: the server keeps the log socket open
        (heartbeats only) after a job ends, so we poll the job state to detect completion.

        Args:
            job_id: The unique identifier of the job.
            teamspace_id: The ID of the teamspace that owns the job.

        Returns:
            ``True`` if the job is Completed/Failed/Stopped, else ``False`` (including when the
            state cannot be determined, so a transient status-check failure never kills a stream).
        """
        from lightning_sdk.status import Status

        try:
            job = self.get_job(job_id, teamspace_id)
        except Exception:
            # a status-check failure must not kill an active stream
            return False
        return self._job_state_to_external(job.state) in (Status.Stopped, Status.Completed, Status.Failed)

    def _job_state_to_external(self, state: str) -> "Status":
        """Convert a raw v2 job state string to the public ``Status`` enum.

        Args:
            state: The raw state string returned by the v2 jobs service.

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

        states = {
            self.v2_job_state_pending: Status.Pending,
            self.v2_job_state_creating: Status.Pending,
            self.v2_job_state_running: Status.Running,
            self.v2_job_state_restarting: Status.Running,
            self.v2_job_state_stopping: Status.Stopping,
            self.v2_job_state_stop: Status.Stopping,
            self.v2_job_state_delete: Status.Stopping,
            # The workload has exited, but its logs and artifacts aren't final until the backend
            # settles the job as completed or failed. Stopping also keeps stop_job() from
            # requesting a stop, which would relabel the job as stopped.
            self.v2_job_state_complete: Status.Stopping,
            self.v2_job_state_fail: Status.Stopping,
            self.v2_job_state_stopped: Status.Stopped,
            self.v2_job_state_deleted: Status.Stopped,
            self.v2_job_state_completed: Status.Completed,
            self.v2_job_state_failed: Status.Failed,
        }
        return states.get(state, Status.Pending)

    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 warn_if_max_runtime_noop(
        self,
        *,
        max_runtime: Optional[int],
        machine: Union[Machine, str],
        interruptible: bool,
        teamspace_id: str,
        cloud_account_id: Optional[str],
        org_id: str,
        stacklevel: int = 3,
    ) -> None:
        """Warn when ``max_runtime`` will not take effect, or will not stop the job.

        ``max_runtime`` maps to a DWS reservation duration. The product UI only sends it for
        non-spot machines with ``dws_supported`` / ``dws_only``; elsewhere it is a no-op. On
        Lightning baremetal (``MACHINE``) it reserves the machine but never stops the job.
        """
        if not max_runtime:
            return

        if interruptible:
            warnings.warn(
                "max_runtime has no effect on interruptible (spot) instances; "
                "it only applies to DWS (Dynamic Workload Scheduler) machines.",
                UserWarning,
                stacklevel=stacklevel,
            )
            return

        if not cloud_account_id:
            return

        try:
            accelerators = self._get_machines_for_cloud_account(
                teamspace_id=teamspace_id,
                cloud_account_id=cloud_account_id,
                org_id=org_id,
            )
        except Exception:
            return

        accelerator = _match_accelerator(machine, accelerators)
        if accelerator is None:
            return
        if str(accelerator.provider) == V1CloudProvider.MACHINE:
            warnings.warn(
                f"max_runtime does not stop the job on machine '{machine}'. On Lightning baremetal it is the "
                "minimum time the machine is reserved for the job; afterwards the job keeps running and the "
                "machine becomes interruptible. For a hard time limit, wrap your command, "
                f"e.g. 'timeout {max_runtime} <command>'.",
                UserWarning,
                stacklevel=stacklevel,
            )
            return
        if accelerator.dws_supported or accelerator.dws_only:
            return

        warnings.warn(
            f"max_runtime has no effect on machine '{machine}'; "
            "it only applies to DWS (Dynamic Workload Scheduler) machines.",
            UserWarning,
            stacklevel=stacklevel,
        )

    @staticmethod
    @lru_cache(maxsize=None)
    def _get_machines_for_cloud_account(
        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.
        """
        from lightning_sdk.api.cloud_account_api import CloudAccountApi

        cloud_account_api = CloudAccountApi()
        accelerators = cloud_account_api.list_cloud_account_accelerators(
            teamspace_id=teamspace_id,
            cloud_account_id=cloud_account_id,
            org_id=org_id,
        )
        if not accelerators:
            return []

        return list(filter(lambda acc: acc.enabled, accelerators.accelerator))

    def get_total_cost(self, job: V1Job) -> float:
        """Return the accumulated cost for a v2 job in USD.

        Args:
            job: The v2 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: V1Job) -> int:
        """Return the number of machines for this job.

        Single-machine jobs always report 1. Kept for parity with ``MMTApiV2.get_num_machines``.
        """
        return 1

    def list_mmt_subjobs(self, job_id: str, teamspace_id: str) -> List[V1Job]:
        """Single-machine jobs have no sub-jobs. Kept for parity with ``MMTApiV2.list_mmt_subjobs``."""
        return []


def _match_accelerator(
    machine: Union[Machine, str],
    accelerators: List[V1ClusterAccelerator],
) -> Optional[V1ClusterAccelerator]:
    """Return the accelerator record matching ``machine``, if any."""
    instance_name = _machine_to_compute_name(machine)
    for accelerator in accelerators:
        possible_identifiers = (
            accelerator.slug,
            accelerator.slug_multi_cloud,
            accelerator.instance_id,
            getattr(accelerator, "secondary_instance_id", None),
        )
        if instance_name in possible_identifiers:
            return accelerator
    return None