from __future__ import annotations

import asyncio
from typing import Any
from itertools import count
from json import JSONDecodeError

from socket import AF_INET

from yarl import URL

from aiohttp import (
    ClientConnectorError,
    ClientOSError,
    ClientSession,
    ClientTimeout,
    ConnectionTimeoutError,
    ContentTypeError,
    TCPConnector,
    ClientResponse,
    BasicAuth,
)

import pwd
import os

from flux_networking_shared.log import log


class ExecBinaryError(ChildProcessError):
    def __init__(self, cmd: list[str], stdout: bytes, stderr: bytes):
        self.__cmd = cmd
        self.__stderr = stderr
        super().__init__([cmd, (stdout, stderr)])

    def stderr(self):
        return self.__stderr.decode(errors="ignore")

    def __str__(self):
        return "Failed to execute '{}', error: {}".format(
            " ".join(self.__cmd), self.stderr()
        )


def merge_dict(a: dict, b: dict, path=[]):
    """Recursively merge dicts (opinionated)

    Args:
        a (dict): The mutated dict dict
        b (dict): The overriding dict
        path (list, optional): where in the chain we are at. Defaults to [].

    Returns:
        _type_: The merged dict. dict b values take precedence.
    """
    for key in b:
        if key in a:
            if isinstance(a[key], dict) and isinstance(b[key], dict):
                merge_dict(a[key], b[key], path + [str(key)])
            else:
                a[key] = b[key]
        else:
            a[key] = b[key]
    return a


def demote(user_uid, user_gid):
    def result():
        os.setgid(user_gid)
        os.setuid(user_uid)

    return result


async def exec_binary(
    cmd: list[str],
    expect_returncode: int = 0,
    cwd: str | None = None,
    user: str | None = None,
    env: dict | None = None,
) -> str:
    if len(cmd) == 0:
        raise ChildProcessError("Cannot execute empty cmd")

    extra_params: dict[str, Any] = {}

    if cwd is not None:
        extra_params["cwd"] = cwd

    if user is not None:
        try:
            pw_record = pwd.getpwnam(user)
        except KeyError:
            raise ChildProcessError("User not found")

        extra_params["preexec_fn"] = demote(pw_record.pw_uid, pw_record.pw_gid)

    if env:
        extra_params["env"] = env

    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            **extra_params,
        )
    except FileNotFoundError:
        raise ChildProcessError(f"Binary not found: {cmd[0]}, is it installed?")

    try:
        stdout, stderr = await proc.communicate()
        if proc.returncode != expect_returncode:
            raise ExecBinaryError(cmd, stdout, stderr)
    except asyncio.CancelledError:
        if proc.returncode is None:
            proc.terminate()
            return_code = await proc.wait()
            log.info(
                "Cmd: %s cancelled and terminated. Exit code: %s",
                cmd[0],
                return_code,
            )
        raise asyncio.CancelledError()

    return stdout.decode()


async def do_http(
    url: URL | str,
    verb: str = "get",
    *,
    data: Any = None,
    connect_timeout: int = 3,
    max_tries: int = 3,
    retry_interval: int = 1,
    credentials: dict | None = None,
    headers: dict | None = None,
    verify_ssl: bool | None = None,
    total_timeout: int | None = None,
) -> Any:
    timeout = ClientTimeout(connect=connect_timeout, total=total_timeout)
    conn = TCPConnector(family=AF_INET)

    auth = None
    res = None

    if credentials:
        # user password
        auth = BasicAuth(*credentials)

    async with ClientSession(
        connector=conn,
    ) as session:
        for attempt in count(start=1):
            if max_tries and attempt > max_tries:
                break

            remaining = max(0, max_tries - attempt)

            try:
                method = getattr(session, verb)
                async with method(
                    url,
                    timeout=timeout,
                    auth=auth,
                    headers=headers,
                    json=data,
                    ssl=verify_ssl,
                ) as resp:
                    resp: ClientResponse  # type: ignore

                    if resp.status in [429, 500, 502, 503, 504]:
                        log.debug("bad response: %s", resp.status)
                        if remaining or not max_tries:
                            await asyncio.sleep(retry_interval)

                        continue

                    if resp.status in [401, 403, 404]:
                        log.info("URL: %s, status: %s", url, resp.status)
                        return None

                    try:
                        res = await resp.json()
                    except (JSONDecodeError, ContentTypeError):
                        res = None
                    break

            except (
                ClientConnectorError,
                ConnectionTimeoutError,
                ClientOSError,
                asyncio.TimeoutError,
            ):
                # ClientOSError: have seen Connection reset by peer
                log.info(
                    "Unable to connect to: %s. Retries: %s remaining. Interval: %ss",
                    url,
                    remaining,
                    retry_interval,
                )

                if remaining or not max_tries:
                    await asyncio.sleep(retry_interval)

    return res
