from __future__ import annotations

import contextlib
import io
import os
import runpy
import subprocess
import sys
from pathlib import Path
from typing import Callable, Iterable, Optional, Tuple


@contextlib.contextmanager
def pushd(path: Path):
    current = Path.cwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(current)


class StreamCapture:
    def __init__(self, on_output: Optional[Callable[[str], None]], collect: io.StringIO) -> None:
        self.on_output = on_output
        self.collect = collect
        self.buffer = ""

    def write(self, text: str) -> int:
        if text is None:
            return 0
        text = str(text)
        self.collect.write(text)
        if self.on_output:
            text = text.replace("\r", "\n")
            self.buffer += text
            while "\n" in self.buffer:
                line, self.buffer = self.buffer.split("\n", 1)
                if line:
                    self.on_output(line)
        return len(text)

    def flush(self) -> None:
        if self.on_output and self.buffer:
            self.on_output(self.buffer)
            self.buffer = ""

    @property
    def encoding(self) -> str:
        return "utf-8"

    @property
    def errors(self) -> str:
        return "replace"

    def isatty(self) -> bool:
        return False

    def writable(self) -> bool:
        return True

    def fileno(self) -> int:
        raise OSError("StreamCapture has no file descriptor")


def run_tool_module(
    tool: str,
    args: Iterable[str],
    cwd: Path,
    stream: bool,
    on_output: Optional[Callable[[str], None]],
) -> Tuple[int, str]:
    argv_backup = sys.argv[:]
    prev_rich_disable = os.environ.get("RICH_DISABLE")
    prev_rich_click_disable = os.environ.get("RICH_CLICK_DISABLE")
    output = io.StringIO()
    capture = StreamCapture(on_output if stream else None, output)
    try:
        sys.argv = [tool, *list(args)]
        os.environ["RICH_DISABLE"] = "1"
        os.environ["RICH_CLICK_DISABLE"] = "1"
        with pushd(cwd), contextlib.redirect_stdout(capture), contextlib.redirect_stderr(capture):
            try:
                module_name = f"{tool}.__main__"
                try:
                    runpy.run_module(module_name, run_name="__main__")
                except ModuleNotFoundError as exc:
                    if exc.name in (module_name, tool):
                        runpy.run_module(tool, run_name="__main__")
                    else:
                        raise
                code = 0
            except SystemExit as exc:
                code = int(exc.code) if isinstance(exc.code, int) else 0
            except Exception:
                import traceback

                traceback.print_exc()
                code = 1
    finally:
        capture.flush()
        sys.argv = argv_backup
        if prev_rich_disable is None:
            os.environ.pop("RICH_DISABLE", None)
        else:
            os.environ["RICH_DISABLE"] = prev_rich_disable
        if prev_rich_click_disable is None:
            os.environ.pop("RICH_CLICK_DISABLE", None)
        else:
            os.environ["RICH_CLICK_DISABLE"] = prev_rich_click_disable
    return code, output.getvalue()


def run_esptool_api(
    func: Callable[[], None],
    stream: bool,
    on_output: Optional[Callable[[str], None]],
) -> Tuple[int, str]:
    output = io.StringIO()
    capture = StreamCapture(on_output if stream else None, output)
    try:
        with contextlib.redirect_stdout(capture), contextlib.redirect_stderr(capture):
            func()
        code = 0
    except SystemExit as exc:
        code = int(exc.code) if isinstance(exc.code, int) else 1
    except Exception:
        import traceback

        traceback.print_exc()
        code = 1
    finally:
        capture.flush()
    return code, output.getvalue()


def run_command(
    cmd: Iterable[str],
    cwd: Path,
    stream: bool = False,
    on_output: Optional[Callable[[str], None]] = None,
) -> Tuple[int, str]:
    if stream:
        proc = subprocess.Popen(
            list(cmd),
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
        output_parts = []
        assert proc.stdout is not None
        for line in proc.stdout:
            output_parts.append(line)
            if on_output:
                on_output(line.rstrip("\n"))
        proc.wait()
        return proc.returncode, "".join(output_parts)

    result = subprocess.run(list(cmd), cwd=cwd, capture_output=True, text=True)
    combined = (result.stdout or "") + (result.stderr or "")
    if on_output and combined:
        for line in combined.splitlines():
            on_output(line)
    return result.returncode, combined