# Copyright 2022 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import logging
import os
import re
import socket
import time
from jax._src import clusters
from jax._src.cloud_tpu_init import running_in_cloud_tpu_vm

logger = logging.getLogger(__name__)

# We use an arbitrarily chosen port for the coordinator since we cannot
# rely on communication to choose one in real time.
coordinator_port = '8482'

metadata_response_code_success = 200

def get_metadata(key):
  import requests  # pytype: disable=import-error
  import time  # pytype: disable=import-error
  # Based on https://github.com/tensorflow/tensorflow/pull/40317
  gce_metadata_endpoint = 'http://' + os.environ.get(
      'GCE_METADATA_IP', 'metadata.google.internal')

  retry_count = 0
  retrySeconds = 0.500
  api_resp = None

  while retry_count < 6:
    api_resp = requests.get(
        f'{gce_metadata_endpoint}/computeMetadata/v1/instance/attributes/{key}',
        headers={'Metadata-Flavor': 'Google'}, timeout=60)
    if api_resp.status_code == 200:
      break
    retry_count += 1
    time.sleep(retrySeconds)

  if api_resp is None:
    raise RuntimeError(f"Getting metadata['{key}'] failed for 6 tries")
  return api_resp.text, api_resp.status_code

def get_tpu_env_value_from_metadata(key) -> str | None:
  metadata_value = None
  tpu_env_data = get_metadata('tpu-env')[0]
  key_value_pairs = tpu_env_data.split('\n')
  for key_value_pair in key_value_pairs:
    # Typical line is MEGASCALE_NUM_SLICES: '2'
    if ':' in key_value_pair:
      row_key, value = re.split(':', key_value_pair, 1)
      row_key = row_key.strip()
      if row_key == key:
        metadata_value = value.strip().strip("'")
  return metadata_value

def get_tpu_env_value(key) -> str | None:
  # First try to get the value from the environment.
  value = os.environ.get(key, None)
  if value is None:
    # If not found, try to get it from the metadata.
    value = get_tpu_env_value_from_metadata(key)
  return value

class BaseTpuCluster(clusters.ClusterEnv):

  name: str = "tpu"

  """Abstract cluster supports both single and multislice TPU environments.

  If MEGASCALE_COORDINATOR_ADDRESS is not set, we assume single slice topology.
  Concrete extensions of this class must implement methods for generating a list
    of within-slice workers and a within-slice process ID.
  `get_coordinator_address` must return the address of the host with
  process ID 0 (as returned by `get_process_id`), since the coordinator service
  is started on the host with process ID = 0.
  """

  @classmethod
  def is_env_present(cls) -> bool:
    """Override this method to return True if the environment is present."""
    return False

  @classmethod
  def get_coordinator_address(cls, timeout_secs: int | None, override_coordinator_port: str | None) -> str:
    # For both GCE via QueuedResources and GKE via JobSet, the
    # Megascale coordinator address is set as the host with process id = 0,
    # so can be used as the jax distributed system coordinator.
    coordinator_address = get_tpu_env_value('MEGASCALE_COORDINATOR_ADDRESS')
    if not coordinator_address:
      # For both GCE (QueuedResources and TPUVM create) and GKE via Job API,
      # the workers lists are sorted by process ID so the first one can
      # be used as the jax distributed system coordinator.
      coordinator_address = cls._get_worker_list_in_slice()[0]
    coordinator_address = coordinator_address.split(':')[0]
    logger.debug("TPU Cluster using coordinator address: %s", coordinator_address)
    cls.wait_for_coordinator(coordinator_address, timeout_secs)
    port = override_coordinator_port or coordinator_port
    return f'{coordinator_address}:{port}'

  @classmethod
  def wait_for_coordinator(cls, coordinator_address, timeout_secs):
    # The coordinator may not be up before the other hosts try to
    # communicate with it. We check for its existence with retries.
    coordinator_found = False
    max_time = time.time() + timeout_secs
    coordinator_retry_secs = 5
    while not coordinator_found and time.time() < max_time:
      try:
        ip_address = socket.gethostbyname(coordinator_address)
        coordinator_found = True
        logger.debug("Found coordinator with address %s", coordinator_address)
      except socket.gaierror:
        logger.debug(
            "Failed to recognize coordinator address %s"
            " retrying...", coordinator_address
        )
        time.sleep(coordinator_retry_secs)
    if not coordinator_found:
      raise RuntimeError(f"Failed to recognize coordinator address {coordinator_address}")

  @classmethod
  def get_process_count(cls) -> int:
    processes_per_slice = len(cls._get_worker_list_in_slice())
    num_slices = cls._get_num_slices()
    total_process_count = processes_per_slice * num_slices
    logger.debug("Total process count of %s = %s processes per slice and %s slices", total_process_count, processes_per_slice, num_slices)
    return total_process_count

  @classmethod
  def get_process_id(cls) -> int:
    process_id_in_slice = cls._get_process_id_in_slice()
    slice_id = cls._get_slice_id()
    processes_per_slice = len(cls._get_worker_list_in_slice())
    process_id = process_id_in_slice + slice_id * processes_per_slice
    logger.debug("Process ID of %s generated by within-slice id %s and slice id %s", process_id, process_id_in_slice, slice_id)
    return process_id

  @staticmethod
  def _get_num_slices() -> int:
    num_slices = get_tpu_env_value('MEGASCALE_NUM_SLICES')
    if not num_slices:
      return 1
    return int(num_slices)  # type: ignore


  @staticmethod
  def _get_slice_id() -> int:
    slice_id = get_tpu_env_value('MEGASCALE_SLICE_ID')
    if not slice_id:
      return 0
    return int(slice_id)  # type: ignore

  @staticmethod
  def _get_process_id_in_slice() -> int:
    """Returns a process ID that is unique within slice."""
    raise NotImplementedError()

  @staticmethod
  def _get_worker_list_in_slice() -> list[str]:
    """Returns a list of worker endpoints/hostnames within slice."""
    raise NotImplementedError()

class GceTpuCluster(BaseTpuCluster):

  name: str = "gcetpu"

  @classmethod
  def is_env_present(cls) -> bool:
    if not running_in_cloud_tpu_vm:
      logger.debug("Did not detect cloud TPU VM")
      return False
    if os.environ.get("TPU_SKIP_MDS_QUERY") is not None:
      logger.debug("TPU_SKIP_MDS_QUERY is set to True, so it's probably not a GCE TPU cluster.")
      return False
    metadata_response, metadata_code = get_metadata('agent-worker-number')
    if metadata_code == metadata_response_code_success:
      logger.debug("Gce Tpu Cluster detected for Jax Distributed System")
      return True
    else:
      logger.debug("Did not detect Gce Tpu Cluster since agent-worker-number is not set in metadata")
      logger.debug("Metadata code: %s", metadata_code)
      logger.debug("Metadata response: %s", metadata_response)
      return False

  @staticmethod
  def _get_process_id_in_slice() -> int:
    return int(get_metadata('agent-worker-number')[0])

  @staticmethod
  def _get_worker_list_in_slice() -> list[str]:
    workers = get_metadata('worker-network-endpoints')[0].split(',')
    return [worker.split(':')[2] for worker in workers]

class GkeTpuCluster(BaseTpuCluster):

  name: str = "gketpu"

  @classmethod
  def is_env_present(cls) -> bool:
    if running_in_cloud_tpu_vm and cls._get_worker_host_names_env_var() is not None:
      logger.debug("Gke Tpu Cluster detected for Jax Distributed System")
      return True
    else:
      if not running_in_cloud_tpu_vm:
        logger.debug("Did not detect cloud TPU VM")
      else:
        logger.debug("Did not detect TPU GKE cluster since neither "
                     "TPU_PROCESS_ADDRESSES nor TPU_WORKER_HOSTNAMES is set.")
      return False

  @staticmethod
  def _get_process_id_in_slice() -> int:
    return int(str(os.environ.get('TPU_WORKER_ID')))

  @staticmethod
  def _get_worker_host_names_env_var() -> str | None:
    """
    Retrieves the list of worker hostnames from environment variables.

    Checks 'TPU_PROCESS_ADDRESSES' first, then 'TPU_WORKER_HOSTNAMES'.
    Returns None if neither environment variable is set.
    """
    worker_hostnames = os.environ.get('TPU_PROCESS_ADDRESSES', None)
    if worker_hostnames is not None:
      return worker_hostnames
    return os.environ.get('TPU_WORKER_HOSTNAMES', None)

  @staticmethod
  def _get_worker_list_in_slice() -> list[str]:
    """
    Returns a list of worker endpoints/hostnames within slice.
    """
    worker_hostnames_str = str(GkeTpuCluster._get_worker_host_names_env_var())
    return worker_hostnames_str.split(',')
