import asyncio
from asyncio.subprocess import SubprocessStreamProtocol
import functools
import os
import pathlib
import subprocess

from typing import Dict, List, Optional

from ceph_devstack import logger, VERBOSE


class LoggingStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
    def __init__(self, limit, loop, log_level):
        self.log_level = log_level
        super().__init__(limit=limit, loop=loop)

    def pipe_data_received(self, fd, data):
        logger.log(
            self.log_level,
            (data.decode() if isinstance(data, bytes) else str(data)).rstrip("\n"),
        )
        super().pipe_data_received(fd, data)


class Command:
    def __init__(
        self,
        args: List[str],
        cwd: Optional[pathlib.Path] = None,
        env: Optional[Dict] = None,
        stream_output: bool = False,
    ):
        self.args = args
        self.env = os.environ | (env or {})
        self.kwargs: Dict = {
            "stdout": asyncio.subprocess.PIPE,
            "stderr": asyncio.subprocess.PIPE,
        }
        if cwd:
            self.kwargs.update(cwd=cwd)
        self.stream_output = stream_output

    def _make_log_msg(self) -> str:
        msg = "> " + " ".join(self.args)
        if (cwd := str(self.kwargs.get("cwd", "."))) != ".":
            msg = f"{msg} cwd='{cwd}'"
        return msg

    def run(self) -> subprocess.Popen:
        logger.log(VERBOSE, self._make_log_msg())
        proc = subprocess.Popen(
            args=self.args,
            env=self.env,
            **self.kwargs,
        )
        proc.wait()
        return proc

    async def arun(self) -> asyncio.subprocess.Process:
        logger.log(VERBOSE, self._make_log_msg())
        loop = asyncio.get_running_loop()
        protocol_factory: functools.partial[SubprocessStreamProtocol]
        if self.stream_output:
            protocol_factory = functools.partial(
                LoggingStreamProtocol,
                limit=2**16,
                loop=loop,
                log_level=VERBOSE,
            )
        else:
            protocol_factory = functools.partial(
                asyncio.subprocess.SubprocessStreamProtocol,
                limit=2**16,
                loop=loop,
            )
        transport, protocol = await loop.subprocess_exec(
            protocol_factory,
            *self.args,
            env=self.env,
            **self.kwargs,
        )
        return asyncio.subprocess.Process(
            transport,
            protocol,
            loop,
        )

    def __str__(self):
        return " ".join(self.args)