# SPDX-License-Identifier: Apache-2.0.
# Copyright (c) 2024 - 2025 Waldiez and contributors.
"""Common utilities for the waldiez runner."""

import asyncio
import os
import re
import subprocess
import sys

# noinspection PyProtectedMember
from asyncio.subprocess import Process
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import AsyncIterator, Iterator, Union


@dataclass
class ProcessSetup:
    """Container for subprocess setup data."""

    temp_dir: Path
    file_path: Path
    old_vars: dict[str, str]
    skip_mmd: bool


@contextmanager
def chdir(to: Union[str, Path]) -> Iterator[None]:
    """Change the current working directory in a context.

    Parameters
    ----------
    to : Union[str, Path]
        The directory to change to.

    Yields
    ------
    Iterator[None]
        The context manager.
    """
    old_cwd = str(os.getcwd())
    os.chdir(to)
    try:
        yield
    finally:
        os.chdir(old_cwd)


@asynccontextmanager
async def a_chdir(to: Union[str, Path]) -> AsyncIterator[None]:
    """Asynchronously change the current working directory in a context.

    Parameters
    ----------
    to : Union[str, Path]
        The directory to change to.

    Yields
    ------
    AsyncIterator[None]
        The async context manager.
    """
    old_cwd = str(os.getcwd())
    os.chdir(to)
    try:
        yield
    finally:
        os.chdir(old_cwd)


def strip_ansi(text: str) -> str:
    """Remove ANSI escape sequences from text.

    Parameters
    ----------
    text : str
        The text to strip.

    Returns
    -------
    str
        The text without ANSI escape sequences.
    """
    ansi_pattern = re.compile(r"\x1b\[[0-9;]*m|\x1b\[.*?[@-~]")
    return ansi_pattern.sub("", text)


def create_sync_subprocess(setup: ProcessSetup) -> subprocess.Popen[bytes]:
    """Create a synchronous subprocess.

    Parameters
    ----------
    setup : ProcessSetup
        The setup data for the subprocess.

    Returns
    -------
    subprocess.Popen[bytes]
        The created subprocess.
    """
    return subprocess.Popen(
        [sys.executable, "-u", str(setup.file_path)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        stdin=subprocess.PIPE,
        # text=True,
        # bufsize=1,  # Line buffered for real-time output
        # universal_newlines=True,
        env={**os.environ},
    )


async def create_async_subprocess(setup: ProcessSetup) -> Process:
    """Create an asynchronous subprocess.

    Parameters
    ----------
    setup : ProcessSetup
        The setup data for the subprocess.

    Returns
    -------
    Process
        The created asynchronous subprocess.
    """
    return await asyncio.create_subprocess_exec(
        sys.executable,
        "-u",
        str(setup.file_path),
        # stdout=asyncio.subprocess.PIPE,
        # stderr=asyncio.subprocess.PIPE,
        # stdin=asyncio.subprocess.PIPE,
        env={**os.environ},
    )
