import fnmatch
import os
import random
import string
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, List, Optional, Union

from lightning_cloud.openapi import (
    Body3,
    Body4,
    Body7,
    Body8,
    Body9,
    Gridv1ImageSpec,
    V1BuildSpec,
    V1DependencyFileInfo,
    V1Drive,
    V1DriveSpec,
    V1DriveStatus,
    V1DriveType,
    V1EnvVar,
    V1Flowserver,
    V1LightningappInstanceSpec,
    V1LightningappInstanceState,
    V1LightningworkDrives,
    V1LightningworkSpec,
    V1Membership,
    V1Metadata,
    V1NetworkConfig,
    V1PackageManager,
    V1ProjectClusterBinding,
    V1PythonDependencyInfo,
    V1QueueServerType,
    V1SourceType,
    V1UserRequestedComputeConfig,
    V1Work,
)
from lightning_cloud.openapi.rest import ApiException

from lightning_app.core.app import LightningApp
from lightning_app.core.constants import (
    CLOUD_QUEUE_TYPE,
    CLOUD_UPLOAD_WARNING,
    DEFAULT_NUMBER_OF_EXPOSED_PORTS,
    DISABLE_DEPENDENCY_CACHE,
    ENABLE_MULTIPLE_WORKS_IN_DEFAULT_CONTAINER,
    ENABLE_MULTIPLE_WORKS_IN_NON_DEFAULT_CONTAINER,
)
from lightning_app.runners.backends.cloud import CloudBackend
from lightning_app.runners.runtime import Runtime
from lightning_app.source_code import LocalSourceCodeDir
from lightning_app.storage import Drive, Mount
from lightning_app.utilities.app_helpers import Logger
from lightning_app.utilities.cloud import _get_project
from lightning_app.utilities.dependency_caching import get_hash
from lightning_app.utilities.packaging.app_config import AppConfig, find_config_file
from lightning_app.utilities.packaging.lightning_utils import _prepare_lightning_wheels_and_requirements
from lightning_app.utilities.secrets import _names_to_ids

logger = Logger(__name__)


@dataclass
class CloudRuntime(Runtime):

    backend: Union[str, CloudBackend] = "cloud"

    def dispatch(
        self,
        on_before_run: Optional[Callable] = None,
        name: str = "",
        cluster_id: str = None,
        **kwargs: Any,
    ) -> None:
        """Method to dispatch and run the :class:`~lightning_app.core.app.LightningApp` in the cloud."""
        # not user facing error ideally - this should never happen in normal user workflow
        if not self.entrypoint_file:
            raise ValueError(
                "Entrypoint file not provided. Did you forget to "
                "initialize the Runtime object with `entrypoint_file` argument?"
            )

        # Determine the root of the project: Start at the entrypoint_file and look for nearby Lightning config files,
        # going up the directory structure. The root of the project is where the Lightning config file is located.

        # TODO: verify lightning version
        # _verify_lightning_version()
        config_file = find_config_file(self.entrypoint_file)
        app_config = AppConfig.load_from_file(config_file) if config_file else AppConfig()
        root = config_file.parent if config_file else Path(self.entrypoint_file).absolute().parent
        cleanup_handle = _prepare_lightning_wheels_and_requirements(root)
        repo = LocalSourceCodeDir(path=root)
        self._check_uploaded_folder(root, repo)
        requirements_file = root / "requirements.txt"
        # The entry point file needs to be relative to the root of the uploaded source file directory,
        # because the backend will invoke the lightning commands relative said source directory
        app_entrypoint_file = Path(self.entrypoint_file).absolute().relative_to(root)

        if name:
            # Override the name if provided by the CLI
            app_config.name = name

        app_config.save_to_dir(root)

        print(f"The name of the app is: {app_config.name}")

        v1_env_vars = [V1EnvVar(name=k, value=v) for k, v in self.env_vars.items()]

        if len(self.secrets.values()) > 0:
            secret_names_to_ids = _names_to_ids(self.secrets.values())
            env_vars_from_secrets = [
                V1EnvVar(name=k, from_secret=secret_names_to_ids[v]) for k, v in self.secrets.items()
            ]
            v1_env_vars.extend(env_vars_from_secrets)

        if ENABLE_MULTIPLE_WORKS_IN_DEFAULT_CONTAINER:
            v1_env_vars.append(V1EnvVar(name="ENABLE_MULTIPLE_WORKS_IN_DEFAULT_CONTAINER", value="1"))

        if ENABLE_MULTIPLE_WORKS_IN_NON_DEFAULT_CONTAINER:
            v1_env_vars.append(V1EnvVar(name="ENABLE_MULTIPLE_WORKS_IN_NON_DEFAULT_CONTAINER", value="1"))

        work_reqs: List[V1Work] = []
        for flow in self.app.flows:
            for work in flow.works(recurse=False):
                work_requirements = "\n".join(work.cloud_build_config.requirements)
                build_spec = V1BuildSpec(
                    commands=work.cloud_build_config.build_commands(),
                    python_dependencies=V1PythonDependencyInfo(
                        package_manager=V1PackageManager.PIP, packages=work_requirements
                    ),
                    image=work.cloud_build_config.image,
                )
                user_compute_config = V1UserRequestedComputeConfig(
                    name=work.cloud_compute.name,
                    count=1,
                    disk_size=work.cloud_compute.disk_size,
                    shm_size=work.cloud_compute.shm_size,
                )

                drive_specs: List[V1LightningworkDrives] = []
                for drive_attr_name, drive in [
                    (k, getattr(work, k)) for k in work._state if isinstance(getattr(work, k), Drive)
                ]:
                    if drive.protocol == "lit://":
                        drive_type = V1DriveType.NO_MOUNT_S3
                        source_type = V1SourceType.S3
                    else:
                        raise RuntimeError(
                            f"unknown drive protocol `{drive.protocol}`. Please verify this "
                            f"drive type has been configured for use in the cloud dispatcher."
                        )

                    drive_specs.append(
                        V1LightningworkDrives(
                            drive=V1Drive(
                                metadata=V1Metadata(
                                    name=f"{work.name}.{drive_attr_name}",
                                ),
                                spec=V1DriveSpec(
                                    drive_type=drive_type,
                                    source_type=source_type,
                                    source=f"{drive.protocol}{drive.id}",
                                ),
                                status=V1DriveStatus(),
                            ),
                            mount_location=str(drive.root_folder),
                        ),
                    )

                # TODO: Move this to the CloudCompute class and update backend
                if work.cloud_compute.mounts is not None:
                    mounts = work.cloud_compute.mounts
                    if isinstance(mounts, Mount):
                        mounts = [mounts]
                    for mount in mounts:
                        drive_specs.append(
                            _create_mount_drive_spec(
                                work_name=work.name,
                                mount=mount,
                            )
                        )

                random_name = "".join(random.choice(string.ascii_lowercase) for _ in range(5))
                spec = V1LightningworkSpec(
                    build_spec=build_spec,
                    cluster_id=cluster_id,
                    drives=drive_specs,
                    user_requested_compute_config=user_compute_config,
                    network_config=[V1NetworkConfig(name=random_name, port=work.port)],
                )
                work_reqs.append(V1Work(name=work.name, spec=spec))

        # We need to collect a spec for each flow that contains a frontend so that the backend knows
        # for which flows it needs to start servers by invoking the cli (see the serve_frontend() method below)
        frontend_specs: List[V1Flowserver] = []
        for flow_name in self.app.frontends.keys():
            frontend_spec = V1Flowserver(name=flow_name)
            frontend_specs.append(frontend_spec)

        app_spec = V1LightningappInstanceSpec(
            app_entrypoint_file=str(app_entrypoint_file),
            enable_app_server=self.start_server,
            flow_servers=frontend_specs,
            desired_state=V1LightningappInstanceState.RUNNING,
            env=v1_env_vars,
        )

        # if requirements file at the root of the repository is present,
        # we pass just the file name to the backend, so backend can find it in the relative path
        if requirements_file.is_file():
            app_spec.image_spec = Gridv1ImageSpec(
                dependency_file_info=V1DependencyFileInfo(package_manager=V1PackageManager.PIP, path="requirements.txt")
            )
            if not DISABLE_DEPENDENCY_CACHE and not kwargs.get("no_cache"):
                # hash used for caching the dependencies
                app_spec.dependency_cache_key = get_hash(requirements_file)
        # we'll get the default project (quite similar to Github Organization) from the backend
        project = _get_project(self.backend.client)

        try:
            list_apps_resp = self.backend.client.lightningapp_v2_service_list_lightningapps_v2(
                project_id=project.project_id, name=app_config.name
            )
            if list_apps_resp.lightningapps:
                # There can be only one app with unique project_id<>name pair
                lit_app = list_apps_resp.lightningapps[0]
            else:
                app_body = Body7(name=app_config.name, can_download_source_code=True)
                lit_app = self.backend.client.lightningapp_v2_service_create_lightningapp_v2(
                    project_id=project.project_id, body=app_body
                )

            release_body = Body8(
                app_entrypoint_file=app_spec.app_entrypoint_file,
                enable_app_server=app_spec.enable_app_server,
                flow_servers=app_spec.flow_servers,
                image_spec=app_spec.image_spec,
                cluster_id=cluster_id,
                works=[V1Work(name=work_req.name, spec=work_req.spec) for work_req in work_reqs],
                local_source=True,
                dependency_cache_key=app_spec.dependency_cache_key,
            )

            if ENABLE_MULTIPLE_WORKS_IN_DEFAULT_CONTAINER:
                network_configs: List[V1NetworkConfig] = []

                initial_port = 8080 + 1 + len(frontend_specs)
                for _ in range(DEFAULT_NUMBER_OF_EXPOSED_PORTS):
                    network_configs.append(
                        V1NetworkConfig(
                            name="w" + str(initial_port),
                            port=initial_port,
                        )
                    )
                    initial_port += 1

                release_body.network_config = network_configs

            if cluster_id is not None:
                self._ensure_cluster_project_binding(project.project_id, cluster_id)

            lightning_app_release = self.backend.client.lightningapp_v2_service_create_lightningapp_release(
                project_id=project.project_id, app_id=lit_app.id, body=release_body
            )

            if cluster_id is not None:
                logger.info(f"running app on {lightning_app_release.cluster_id}")

            if lightning_app_release.source_upload_url == "":
                raise RuntimeError("The source upload url is empty.")

            repo.package()
            repo.upload(url=lightning_app_release.source_upload_url)

            # check if user has sufficient credits to run an app
            # if so set the desired state to running otherwise, create the app in stopped state,
            # and open the admin ui to add credits and running the app.
            has_sufficient_credits = self._project_has_sufficient_credits(project, app=self.app)
            app_release_desired_state = (
                V1LightningappInstanceState.RUNNING if has_sufficient_credits else V1LightningappInstanceState.STOPPED
            )
            if not has_sufficient_credits:
                logger.warn("You may need Lightning credits to run your apps on the cloud.")

            # right now we only allow a single instance of the app
            find_instances_resp = self.backend.client.lightningapp_instance_service_list_lightningapp_instances(
                project_id=project.project_id, app_id=lit_app.id
            )
            queue_server_type = V1QueueServerType.REDIS if CLOUD_QUEUE_TYPE == "redis" else V1QueueServerType.HTTP
            if find_instances_resp.lightningapps:
                existing_instance = find_instances_resp.lightningapps[0]
                if existing_instance.status.phase != V1LightningappInstanceState.STOPPED:
                    # TODO(yurij): Implement release switching in the UI and remove this
                    # We can only switch release of the stopped instance
                    existing_instance = self.backend.client.lightningapp_instance_service_update_lightningapp_instance(
                        project_id=project.project_id,
                        id=existing_instance.id,
                        body=Body3(spec=V1LightningappInstanceSpec(desired_state=V1LightningappInstanceState.STOPPED)),
                    )
                    # wait for the instance to stop for up to 150 seconds
                    for _ in range(150):
                        existing_instance = self.backend.client.lightningapp_instance_service_get_lightningapp_instance(
                            project_id=project.project_id, id=existing_instance.id
                        )
                        if existing_instance.status.phase == V1LightningappInstanceState.STOPPED:
                            break
                        time.sleep(1)
                    if existing_instance.status.phase != V1LightningappInstanceState.STOPPED:
                        raise RuntimeError("Failed to stop the existing instance.")

                lightning_app_instance = (
                    self.backend.client.lightningapp_instance_service_update_lightningapp_instance_release(
                        project_id=project.project_id,
                        id=existing_instance.id,
                        body=Body4(release_id=lightning_app_release.id),
                    )
                )

                self.backend.client.lightningapp_instance_service_update_lightningapp_instance(
                    project_id=project.project_id,
                    id=existing_instance.id,
                    body=Body3(
                        spec=V1LightningappInstanceSpec(
                            desired_state=app_release_desired_state,
                            env=v1_env_vars,
                            queue_server_type=queue_server_type,
                        )
                    ),
                )
            else:
                lightning_app_instance = (
                    self.backend.client.lightningapp_v2_service_create_lightningapp_release_instance(
                        project_id=project.project_id,
                        app_id=lit_app.id,
                        id=lightning_app_release.id,
                        body=Body9(
                            cluster_id=cluster_id,
                            desired_state=app_release_desired_state,
                            name=lit_app.name,
                            env=v1_env_vars,
                            queue_server_type=queue_server_type,
                        ),
                    )
                )
        except ApiException as e:
            logger.error(e.body)
            sys.exit(1)

        if on_before_run:
            on_before_run(lightning_app_instance, need_credits=not has_sufficient_credits)

        if lightning_app_instance.status.phase == V1LightningappInstanceState.FAILED:
            raise RuntimeError("Failed to create the application. Cannot upload the source code.")

        if cleanup_handle:
            cleanup_handle()

    def _ensure_cluster_project_binding(self, project_id: str, cluster_id: str):
        cluster_bindings = self.backend.client.projects_service_list_project_cluster_bindings(project_id=project_id)

        for cluster_binding in cluster_bindings.clusters:
            if cluster_binding.cluster_id != cluster_id:
                continue
            if cluster_binding.project_id == project_id:
                return

        self.backend.client.projects_service_create_project_cluster_binding(
            project_id=project_id,
            body=V1ProjectClusterBinding(cluster_id=cluster_id, project_id=project_id),
        )

    @staticmethod
    def _check_uploaded_folder(root: Path, repo: LocalSourceCodeDir) -> None:
        """This method is used to inform the users if their folder files are large and how to filter them."""
        lightning_tar = set(fnmatch.filter(repo.files, "*lightning-*.tar.gz"))
        app_folder_size = sum(Path(p).stat().st_size for p in repo.files if p not in lightning_tar)
        app_folder_size_in_mb = round(app_folder_size / (1000 * 1000), 5)
        if app_folder_size_in_mb > CLOUD_UPLOAD_WARNING:
            path_sizes = [(p, Path(p).stat().st_size / (1000 * 1000)) for p in repo.files]
            largest_paths = sorted((x for x in path_sizes if x[-1] > 0.01), key=lambda x: x[1], reverse=True)[:25]
            largest_paths_msg = "\n".join(f"{round(s, 5)} MB: {p}" for p, s in largest_paths)
            warning_msg = (
                f"Your application folder {root} is more than {CLOUD_UPLOAD_WARNING} MB. "
                f"Found {app_folder_size_in_mb} MB \n"
                "Here are the largest files: \n"
                f"{largest_paths_msg}"
            )
            if not os.path.exists(os.path.join(root, ".lightningignore")):
                warning_msg = (
                    warning_msg
                    + "\nIn order to ignore some files or folder, "
                    + "create a `.lightningignore` file and add the paths to ignore."
                )
            else:
                warning_msg += "\nYou can ignore some files or folders by adding them to `.lightningignore`."

            logger.warn(warning_msg)

    def _project_has_sufficient_credits(self, project: V1Membership, app: Optional[LightningApp] = None):
        """check if user has enough credits to run the app with its hardware if app is not passed return True if
        user has 1 or more credits."""
        balance = project.balance
        if balance is None:
            balance = 0  # value is missing in some tests

        return balance >= 1


def _create_mount_drive_spec(work_name: str, mount: Mount) -> V1LightningworkDrives:
    if mount.protocol == "s3://":
        drive_type = V1DriveType.INDEXED_S3
        source_type = V1SourceType.S3
    else:
        raise RuntimeError(
            f"unknown mount protocol `{mount.protocol}`. Please verify this "
            f"drive type has been configured for use in the cloud dispatcher."
        )

    return V1LightningworkDrives(
        drive=V1Drive(
            metadata=V1Metadata(
                name=work_name,
            ),
            spec=V1DriveSpec(
                drive_type=drive_type,
                source_type=source_type,
                source=mount.source,
            ),
            status=V1DriveStatus(),
        ),
        mount_location=str(mount.mount_path),
    )
