'''
'''
#  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.

from contextlib import contextmanager
from typing import Set
import socket
import subprocess
import os
import platform
import psutil

import hosts.output as host

from ordered_set_queue import OrderedSetQueue

g_ports = None  # ports we can use


class PortQueueSelectionError(Exception):
    """
    An exception for when there are problems selecting a port from the port
    queue.
    """
    pass


def PortOpen(port: int, address: str = None, bound_ports: Set[int] = None) -> bool:
    """
    Detect whether the port is open, that is a socket is currently using that port.

    Open ports are currently in use by an open socket and are therefore not available
    for a server to listen on them.

    Args:
        port: The port to check.
        address: The address to check. Defaults to localhost.
        bound_ports: A set of ports that are currently bound. If a port is in
            this set, it is considered open.

    Returns:
        True if a socket is currently bound to the port or accepts a TCP
        connection, False otherwise.
    """
    ret = False
    if address is None:
        address = "localhost"

    if port in bound_ports:
        host.WriteDebug('PortOpen', f"{port} is open because it is in the bound sockets set.")
        return True

    address = (address, port)

    try:
        # Try to connect on that port. If we can connect on it, then someone is
        # listening on that port and therefore the port is open.
        s = socket.create_connection(address, timeout=.5)
        s.close()
        ret = True
        host.WriteDebug(
            'PortOpen', f"Connection to port {port} succeeded, the port is open, "
            "and a future connection cannot use it")
    except OSError:
        host.WriteDebug(
            'PortOpen', f"socket error for port {port}, port is closed, "
            "and therefore a future connection can use it")
    except socket.timeout:
        host.WriteDebug(
            'PortOpen', f"Timeout error for port {port}, port is closed, "
            "and therefore a future connection can use it")

    return ret


def _get_available_port(queue):
    """
    Get the next available port from the port queue and return it.

    Since there can be a delay between when the queue is populated and when the
    port is requested, this checks the next port in the queue to see whether it
    is still not open. If it isn't, it is dropped from the queue and the next
    one is inspected. This process continues until either an available port is
    found, or, if the queue is exhausted, PortQueueSelectionError is raised.

    Returns:
        An available (i.e., non-open) port.

    Throws:
        PortQueueSelectionError if the port queue is exhausted.
    """

    if queue.qsize() == 0:
        host.WriteWarning("Port queue is empty.")
        raise PortQueueSelectionError("Could not get a valid port because the queue is empty")

    bound_ports = _get_bound_ports()
    port = queue.get()
    while PortOpen(port, bound_ports=bound_ports):
        host.WriteDebug('_get_available_port', f"Port was closed but now is used: {port}")
        if queue.qsize() == 0:
            host.WriteWarning("Port queue is empty.")
            raise PortQueueSelectionError("Could not get a valid port because the queue is empty")
        port = queue.get()
    return port


def _is_bound(conn) -> bool:
    """Return whether an internet socket connection occupies its local port."""
    return bool(
        conn.family in (socket.AF_INET, socket.AF_INET6) and conn.laddr and
        (conn.status == psutil.CONN_LISTEN or conn.type == socket.SOCK_DGRAM))


def _get_bound_ports() -> Set[int]:
    """Use psutil to get the set of ports that are currently bound.

    TCP sockets report a listening status, but UDP sockets have no comparable
    status. Any UDP socket with a local address is bound and therefore makes
    its port unavailable to AuTest processes.

    :return: The set of ports that are currently bound.
    """
    ports: Set[int] = set()
    try:
        connections = psutil.net_connections(kind='all')
        for conn in connections:
            if _is_bound(conn):
                ports.add(conn.laddr.port)
    except psutil.AccessDenied:
        # Mac OS X doesn't allow net_connections() to be called without root.
        for proc in psutil.process_iter(['pid', 'name']):
            try:
                connections = proc.connections(kind='all')
            except (psutil.AccessDenied, psutil.NoSuchProcess):
                continue
            for conn in connections:
                if _is_bound(conn):
                    ports.add(conn.laddr.port)
    return ports


def _setup_port_queue(amount=1000):
    """
    Build up the set of ports that the OS in theory will not use.

    The AUTEST_PORT_OFFSET environment variable can be used to offset the
    starting port range. This is useful when running multiple autest processes
    in parallel to avoid port conflicts. Each parallel worker should use a
    different offset (e.g., 0, 1000, 2000, etc.).
    """
    global g_ports
    if g_ports is None:
        host.WriteDebug('_setup_port_queue', "Populating the port queue.")
        g_ports = OrderedSetQueue()
    else:
        # The queue has already been populated.
        host.WriteDebug('_setup_port_queue', f"Queue was previously populated. Queue size: {g_ports.qsize()}")
        return

    # Get port offset for parallel execution support
    try:
        port_offset = int(os.environ.get('AUTEST_PORT_OFFSET', 0))
    except ValueError:
        host.WriteWarning("AUTEST_PORT_OFFSET is not a valid integer, defaulting to 0")
        port_offset = 0
    # Clamp to a safe range to avoid exceeding the valid port space
    port_offset = max(0, min(port_offset, 60000))
    if port_offset > 0:
        host.WriteVerbose('_setup_port_queue', f"Using port offset: {port_offset}")

    try:
        # Use sysctl to find the range of ports that the OS publishes it uses.
        # some docker setups don't have sbin setup correctly
        new_env = os.environ.copy()
        new_env['PATH'] = "/sbin:/usr/sbin:" + new_env['PATH']
        if 'Darwin' == platform.system():
            dmin = subprocess.check_output(["sysctl", "net.inet.ip.portrange.first"], env=new_env).decode().split(":")[1].split()[0]
            dmax = subprocess.check_output(["sysctl", "net.inet.ip.portrange.last"], env=new_env).decode().split(":")[1].split()[0]
        else:
            dmin, dmax = subprocess.check_output(["sysctl", "net.ipv4.ip_local_port_range"],
                                                 env=new_env).decode().split("=")[1].split()
        dmin = int(dmin)
        dmax = int(dmax)
    except Exception:
        host.WriteWarning("Unable to call sysctrl!\n Tests may fail because of bad port selection!")
        return

    rmin = dmin - 2000
    rmax = 65536 - dmax

    bound_ports = _get_bound_ports()
    if rmax > amount:
        # Fill in ports, starting above the upper OS-usable port range.
        # Add port_offset to support parallel test execution.
        port = dmax + 1 + port_offset
        while port < 65536 and g_ports.qsize() < amount:
            if PortOpen(port, bound_ports=bound_ports):
                host.WriteDebug('_setup_port_queue', f"Rejecting an already bound port: {port}")
            else:
                host.WriteDebug('_setup_port_queue', f"Adding a possible port to connect to: {port}")
                g_ports.put(port)
            port += 1
    if rmin > amount and g_ports.qsize() < amount:
        # Fill in more ports, starting at 2001, well above well known ports,
        # and going up until the minimum port range used by the OS.
        # Add port_offset to support parallel test execution (same as high range).
        port = 2001 + port_offset
        while port < dmin and g_ports.qsize() < amount:
            if PortOpen(port, bound_ports=bound_ports):
                host.WriteDebug('_setup_port_queue', f"Rejecting an already bound port: {port}")
            else:
                host.WriteDebug('_setup_port_queue', f"Adding a possible port to connect to: {port}")
                g_ports.put(port)
            port += 1


def _get_port_by_bind():
    """
    Create a socket, bind with REUSEADDR, and return the bound port.

    Because SO_REUSEADDR is applied as an option to the socket, the future
    server should be able to use this port. This method is considered
    sub-optimal in comparison to the OS-unused port range method used above
    because another process might bind to this port in the meantime.

    Returns:
        A port value that can be used for a connection.
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(('', 0))  # bind to all interfaces on an ephemeral port
    port = sock.getsockname()[1]
    return port


def _reserve_port():
    """
    Get a port from the global port queue.

    Returns:
        A tuple containing the port value and whether it should be recycled
        into the queue when the caller is done with it.
    """
    _setup_port_queue()
    if g_ports.qsize() > 0:
        try:
            port = _get_available_port(g_ports)
            host.WriteVerbose("_reserve_port", f"Using port from port queue: {port}")
            return port, True
        except PortQueueSelectionError:
            port = _get_port_by_bind()
            host.WriteVerbose("_reserve_port", f"Queue was drained. Using port from a bound socket: {port}")
            return port, False

    # Since the queue could not be populated, use a port via bind.
    port = _get_port_by_bind()
    host.WriteVerbose("_reserve_port", f"Queue is empty. Using port from a bound socket: {port}")
    return port, False


@contextmanager
def get_port_number():
    """
    Reserve a port number from the same allocator used by get_port().

    This is useful for helper code that needs a temporary listening port but
    does not have an AuTest object with Setup hooks for recycling it. Queue
    ports are recycled when the context exits.

    :returns: A context manager yielding the reserved port value.
    """
    port, recycle_port = _reserve_port()
    try:
        yield port
    finally:
        if recycle_port:
            g_ports.put(port)


def get_port(obj, name):
    '''
    Get a port and set it to the specified variable on the object.

    Args:
        obj: The object upon which to set the port variable.
        name: The name of the variable that receives the port value.

    Returns:
        The port value.
    '''
    port, recycle_port = _reserve_port()
    if recycle_port:
        obj.Setup.Lambda(
            func_cleanup=lambda: g_ports.put(port), description=f"recycling port: {port}, queue size: {g_ports.qsize()}")

    # Assign to the named variable.
    obj.Variables[name] = port
    return port
