import time

from ray_release.anyscale_util import create_cluster_env_from_image
from ray_release.cluster_manager.cluster_manager import ClusterManager
from ray_release.config import COMPUTE_CONFIG_MODEL_FIELDS
from ray_release.exception import (
    ClusterComputeCreateError,
    ClusterEnvBuildError,
    ClusterEnvBuildTimeout,
    ClusterEnvCreateError,
)
from ray_release.logger import logger
from ray_release.retry import retry
from ray_release.util import (
    anyscale_cluster_env_build_url,
    format_link,
)

REPORT_S = 30.0


class MinimalClusterManager(ClusterManager):
    """Minimal manager.

    Builds app config and compute template but does not start or stop session.
    """

    @retry(
        init_delay_sec=10,
        jitter_sec=5,
        max_retry_count=2,
        exceptions=(ClusterEnvCreateError,),
    )
    def create_cluster_env(self):
        assert self.cluster_env_id is None
        assert self.cluster_env_name

        logger.info(
            f"Test uses a cluster env with name "
            f"{self.cluster_env_name}. Looking up existing "
            f"cluster envs with this name."
        )

        self.cluster_env_id = create_cluster_env_from_image(
            image=self.test.get_anyscale_byod_image(),
            test_name=self.cluster_env_name,
            runtime_env=self.test.get_byod_runtime_env(),
            sdk=self.sdk,
            cluster_env_id=self.cluster_env_id,
            cluster_env_name=self.cluster_env_name,
        )

    def build_cluster_env(self, timeout: float = 600.0):
        assert self.cluster_env_id
        assert self.cluster_env_build_id is None

        # Fetch build
        build_id = None
        last_status = None
        error_message = None
        config_json = None
        result = self.sdk.list_cluster_environment_builds(self.cluster_env_id)
        if not result or not result.results:
            raise ClusterEnvBuildError(f"No build found for cluster env: {result}")

        build = sorted(result.results, key=lambda b: b.created_at)[-1]
        build_id = build.id
        last_status = build.status
        error_message = build.error_message
        config_json = build.config_json

        if last_status == "succeeded":
            logger.info(
                f"Link to succeeded cluster env build: "
                f"{format_link(anyscale_cluster_env_build_url(build_id))}"
            )
            self.cluster_env_build_id = build_id
            return

        if last_status == "failed":
            logger.info(f"Previous cluster env build failed: {error_message}")
            logger.info("Starting new cluster env build...")

            # Retry build
            result = self.sdk.create_cluster_environment_build(
                dict(
                    cluster_environment_id=self.cluster_env_id, config_json=config_json
                )
            )
            build_id = result.result.id

            logger.info(
                f"Link to created cluster env build: "
                f"{format_link(anyscale_cluster_env_build_url(build_id))}"
            )

        # Build found but not failed/finished yet
        completed = False
        start_wait = time.time()
        next_report = start_wait + REPORT_S
        timeout_at = time.monotonic() + timeout
        logger.info(f"Waiting for build {build_id} to finish...")
        logger.info(
            f"Track progress here: "
            f"{format_link(anyscale_cluster_env_build_url(build_id))}"
        )
        while not completed:
            now = time.time()
            if now > next_report:
                logger.info(
                    f"... still waiting for build {build_id} to finish "
                    f"({int(now - start_wait)} seconds) ..."
                )
                next_report = next_report + REPORT_S

            result = self.sdk.get_build(build_id)
            build = result.result

            if build.status == "failed":
                raise ClusterEnvBuildError(
                    f"Cluster env build failed. Please see "
                    f"{anyscale_cluster_env_build_url(build_id)} for details. "
                    f"Error message: {build.error_message}"
                )

            if build.status == "succeeded":
                logger.info("Build succeeded.")
                self.cluster_env_build_id = build_id
                return

            completed = build.status not in ["in_progress", "pending"]

            if completed:
                raise ClusterEnvBuildError(
                    f"Unknown build status: {build.status}. Please see "
                    f"{anyscale_cluster_env_build_url(build_id)} for details"
                )

            if time.monotonic() > timeout_at:
                raise ClusterEnvBuildTimeout(
                    f"Time out when building cluster env {self.cluster_env_name}"
                )

            time.sleep(1)

        self.cluster_env_build_id = build_id

    def create_cluster_compute(self, _repeat: bool = True):
        assert self.cluster_compute_id is None

        if self.cluster_compute and self.test.uses_anyscale_sdk_2026():
            return self._create_cluster_compute_new_sdk(_repeat=_repeat)

        if self.cluster_compute:
            assert self.cluster_compute

            logger.info(
                f"Tests uses compute template "
                f"with name {self.cluster_compute_name}. "
                f"Looking up existing cluster computes."
            )

            paging_token = None
            while not self.cluster_compute_id:
                result = self.sdk.search_cluster_computes(
                    dict(
                        project_id=self.project_id,
                        name=dict(equals=self.cluster_compute_name),
                        include_anonymous=True,
                        paging=dict(paging_token=paging_token),
                    )
                )
                paging_token = result.metadata.next_paging_token

                for res in result.results:
                    if res.name == self.cluster_compute_name:
                        self.cluster_compute_id = res.id
                        logger.info(
                            f"Cluster compute already exists "
                            f"with ID {self.cluster_compute_id}"
                        )
                        break

                if not paging_token:
                    break

            if not self.cluster_compute_id:
                logger.info(
                    f"Cluster compute not found. "
                    f"Creating with name {self.cluster_compute_name}."
                )
                try:
                    result = self.sdk.create_cluster_compute(
                        dict(
                            name=self.cluster_compute_name,
                            project_id=self.project_id,
                            config=self.cluster_compute,
                        )
                    )
                    self.cluster_compute_id = result.result.id
                except Exception as e:
                    if _repeat:
                        logger.warning(
                            f"Got exception when trying to create cluster "
                            f"compute: {e}. Sleeping for 10 seconds and then "
                            f"try again once..."
                        )
                        time.sleep(10)
                        return self.create_cluster_compute(_repeat=False)

                    raise ClusterComputeCreateError(
                        "Could not create cluster compute"
                    ) from e

                logger.info(
                    f"Cluster compute template created with "
                    f"name {self.cluster_compute_name} and "
                    f"ID {self.cluster_compute_id}"
                )

    def _create_cluster_compute_new_sdk(self, _repeat: bool = True):
        """Create cluster compute using the anyscale.compute_config API (2026 SDK)."""
        assert self.cluster_compute

        logger.info(
            f"Test uses a compute config (2026 SDK) with name "
            f"{self.cluster_compute_name}. Looking up existing "
            f"compute configs with this name."
        )

        # Check if compute config already exists by name
        anyscale_sdk = self.test.anyscale
        try:
            compute_config_version = anyscale_sdk.compute_config.get(
                self.cluster_compute_name
            )
            self.cluster_compute_id = compute_config_version.id
            logger.info(
                f"Compute config already exists " f"with ID {self.cluster_compute_id}"
            )
            return
        except RuntimeError as e:
            if "not found" not in str(e):
                raise
        logger.info(
            f"Compute config not found. "
            f"Creating with name {self.cluster_compute_name}."
        )

        # Build ComputeConfig from the cluster compute dict, excluding
        # keys like idle_termination_minutes/maximum_uptime_minutes that
        # were added by set_cluster_compute() and are not part of the
        # ComputeConfig model.
        ComputeConfig = anyscale_sdk.compute_config.ComputeConfig
        config_dict = {
            k: v
            for k, v in self.cluster_compute.items()
            if k in COMPUTE_CONFIG_MODEL_FIELDS
        }
        config = ComputeConfig.from_dict(config_dict)

        try:
            full_name = anyscale_sdk.compute_config.create(
                config, name=self.cluster_compute_name
            )
            compute_config_version = anyscale_sdk.compute_config.get(full_name)
            self.cluster_compute_id = compute_config_version.id
        except Exception as e:
            if _repeat:
                logger.warning(
                    f"Got exception when trying to create compute "
                    f"config: {e}. Sleeping for 10 seconds and then "
                    f"try again once..."
                )
                time.sleep(10)
                return self._create_cluster_compute_new_sdk(_repeat=False)

            raise ClusterComputeCreateError("Could not create cluster compute") from e

        logger.info(
            f"Compute config created with "
            f"name {full_name} and "
            f"ID {self.cluster_compute_id}"
        )

    def build_configs(self, timeout: float = 30.0):
        try:
            self.create_cluster_compute()
        except AssertionError as e:
            # If already exists, ignore
            logger.warning(str(e))
        except ClusterComputeCreateError as e:
            raise e
        except Exception as e:
            raise ClusterComputeCreateError(
                f"Unexpected cluster compute build error: {e}"
            ) from e

        try:
            self.create_cluster_env()
        except AssertionError as e:
            # If already exists, ignore
            logger.warning(str(e))
        except ClusterEnvCreateError as e:
            raise e
        except Exception as e:
            raise ClusterEnvCreateError(
                f"Unexpected cluster env create error: {e}"
            ) from e

        try:
            self.build_cluster_env(timeout=timeout)
        except AssertionError as e:
            # If already exists, ignore
            logger.warning(str(e))
        except (ClusterEnvBuildError, ClusterEnvBuildTimeout) as e:
            raise e
        except Exception as e:
            raise ClusterEnvBuildError(
                f"Unexpected cluster env build error: {e}"
            ) from e