#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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
#
#    http://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.
#

"""
Util functions for workers.
"""

from contextlib import contextmanager
import importlib
from inspect import currentframe, getframeinfo
import os
import sys
from typing import Any, Generator, IO, Optional
import warnings

if "SPARK_TESTING" in os.environ:
    assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER", (
        "This module can only be imported in python woker"
    )

# 'resource' is a Unix specific module.
has_resource_module = True
try:
    import resource
except ImportError:
    has_resource_module = False

from pyspark.accumulators import _accumulatorRegistry
from pyspark.util import is_remote_only
from pyspark.errors import PySparkRuntimeError
from pyspark.util import local_connect_and_auth
from pyspark.serializers import (
    read_bool,
    read_int,
    read_long,
    write_int,
    FramedSerializer,
    UTF8Deserializer,
    CPickleSerializer,
)

pickleSer = CPickleSerializer()
utf8_deserializer = UTF8Deserializer()


def add_path(path: str) -> bool:
    # worker can be used, so do not add path multiple times
    if path not in sys.path:
        # overwrite system packages
        sys.path.insert(1, path)
        return True
    return False


def read_command(serializer: FramedSerializer, file: IO) -> Any:
    if not is_remote_only():
        from pyspark.core.broadcast import Broadcast

    command = serializer._read_with_length(file)
    if not is_remote_only() and isinstance(command, Broadcast):
        command = serializer.loads(command.value)
    return command


def check_python_version(infile: IO) -> None:
    """
    Check the Python version between the running process and the one used to serialize the command.
    """
    version = utf8_deserializer.loads(infile)
    worker_version = "%d.%d" % sys.version_info[:2]
    if version != worker_version:
        raise PySparkRuntimeError(
            errorClass="PYTHON_VERSION_MISMATCH",
            messageParameters={
                "worker_version": worker_version,
                "driver_version": str(version),
            },
        )


def setup_memory_limits(memory_limit_mb: int) -> None:
    """
    Sets up the memory limits.

    If memory_limit_mb > 0 and `resource` module is available, sets the memory limit.
    Windows does not support resource limiting and actual resource is not limited on MacOS.
    """
    if memory_limit_mb > 0 and has_resource_module:
        total_memory = resource.RLIMIT_AS
        try:
            soft_limit, hard_limit = resource.getrlimit(total_memory)
            msg = "Current mem limits: {0} of max {1}\n".format(soft_limit, hard_limit)
            print(msg, file=sys.stderr)

            # convert to bytes
            new_limit = memory_limit_mb * 1024 * 1024

            if soft_limit == resource.RLIM_INFINITY or new_limit < soft_limit:
                msg = "Setting mem limits to {0} of max {1}\n".format(new_limit, new_limit)
                print(msg, file=sys.stderr)
                resource.setrlimit(total_memory, (new_limit, new_limit))

        except (resource.error, OSError, ValueError) as e:
            # not all systems support resource limits, so warn instead of failing
            current = currentframe()
            lineno = getframeinfo(current).lineno + 1 if current is not None else 0
            if "__file__" in globals():
                print(
                    warnings.formatwarning(
                        "Failed to set memory limit: {0}".format(e),
                        ResourceWarning,
                        __file__,
                        lineno,
                    ),
                    file=sys.stderr,
                )


def setup_spark_files(infile: IO) -> None:
    """
    Set up Spark files, archives, and pyfiles.
    """
    # fetch name of workdir
    spark_files_dir = utf8_deserializer.loads(infile)

    if not is_remote_only():
        from pyspark.core.files import SparkFiles

        SparkFiles._root_directory = spark_files_dir
        SparkFiles._is_running_on_worker = True

    # fetch names of includes (*.zip and *.egg files) and construct PYTHONPATH
    path_changed = add_path(spark_files_dir)  # *.py files that were added will be copied here
    num_python_includes = read_int(infile)
    for _ in range(num_python_includes):
        filename = utf8_deserializer.loads(infile)
        path_changed = add_path(os.path.join(spark_files_dir, filename)) or path_changed

    if path_changed:
        importlib.invalidate_caches()


def setup_broadcasts(infile: IO) -> None:
    """
    Set up broadcasted variables.
    """
    if not is_remote_only():
        from pyspark.core.broadcast import Broadcast, _broadcastRegistry

    # fetch names and values of broadcast variables
    needs_broadcast_decryption_server = read_bool(infile)
    num_broadcast_variables = read_int(infile)
    if needs_broadcast_decryption_server:
        # read the decrypted data from a server in the jvm
        conn_info = read_int(infile)
        auth_secret = None
        if conn_info == -1:
            conn_info = utf8_deserializer.loads(infile)
        else:
            auth_secret = utf8_deserializer.loads(infile)
        broadcast_sock_file, _ = local_connect_and_auth(conn_info, auth_secret)

    for _ in range(num_broadcast_variables):
        bid = read_long(infile)
        if bid >= 0:
            if needs_broadcast_decryption_server:
                read_bid = read_long(broadcast_sock_file)
                assert read_bid == bid
                _broadcastRegistry[bid] = Broadcast(sock_file=broadcast_sock_file)
            else:
                path = utf8_deserializer.loads(infile)
                _broadcastRegistry[bid] = Broadcast(path=path)

        else:
            bid = -bid - 1
            _broadcastRegistry.pop(bid)

    if needs_broadcast_decryption_server:
        broadcast_sock_file.write(b"1")
        broadcast_sock_file.close()


@contextmanager
def get_sock_file_to_executor(timeout: Optional[int] = -1) -> Generator[IO, None, None]:
    # Read information about how to connect back to the JVM from the environment.
    conn_info = os.environ.get(
        "PYTHON_WORKER_FACTORY_SOCK_PATH", int(os.environ.get("PYTHON_WORKER_FACTORY_PORT", -1))
    )
    auth_secret = os.environ.get("PYTHON_WORKER_FACTORY_SECRET")
    sock_file, sock = local_connect_and_auth(conn_info, auth_secret)
    if timeout is None or timeout > 0:
        sock.settimeout(timeout)
    # TODO: Remove the following two lines and use `Process.pid()` when we drop JDK 8.
    write_int(os.getpid(), sock_file)
    sock_file.flush()
    try:
        yield sock_file
    finally:
        sock_file.close()


def send_accumulator_updates(outfile: IO) -> None:
    """
    Send the accumulator updates back to JVM.
    """
    write_int(len(_accumulatorRegistry), outfile)
    for aid, accum in _accumulatorRegistry.items():
        pickleSer._write_with_length((aid, accum._value), outfile)


class Conf:
    def __init__(self, infile: Optional[IO] = None) -> None:
        self._conf: dict[str, Any] = {}
        if infile is not None:
            self.load(infile)

    def load(self, infile: IO) -> None:
        num_conf = read_int(infile)
        # We do a sanity check here to reduce the possibility to stuck indefinitely
        # due to an invalid messsage. If the numer of configurations is obviously
        # wrong, we just raise an error directly.
        # We hand-pick the configurations to send to the worker so the number should
        # be very small (less than 100).
        if num_conf < 0 or num_conf > 10000:
            raise PySparkRuntimeError(
                errorClass="PROTOCOL_ERROR",
                messageParameters={
                    "failure": f"Invalid number of configurations: {num_conf}",
                },
            )
        for _ in range(num_conf):
            k = utf8_deserializer.loads(infile)
            v = utf8_deserializer.loads(infile)
            self._conf[k] = v

    def get(self, key: str, default: Any = "", *, lower_str: bool = True) -> Any:
        val = self._conf.get(key, default)
        if isinstance(val, str) and lower_str:
            return val.lower()
        return val
