from __future__ import annotations as _annotations

import asyncio
import contextlib
import math
import os
import sys
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Literal

from websockets.asyncio.server import ServerConnection
from websockets.exceptions import ConnectionClosed

from .config import TransportOptions

__all__ = ("CommandOptions", "run_remote_command_connection")


@dataclass(frozen=True, kw_only=True)
class CommandOptions:
    command: tuple[str, ...]
    cwd: str | None = None
    env: Mapping[str, str] | None = None
    stderr_mode: Literal["inherit", "discard"] = "inherit"
    terminate_timeout: float = 5.0

    def __post_init__(self) -> None:
        if not self.command:
            raise ValueError("command must not be empty")
        if not math.isfinite(self.terminate_timeout) or self.terminate_timeout <= 0:
            raise ValueError("terminate_timeout must be a positive finite number")


async def run_remote_command_connection(
    websocket: ServerConnection,
    *,
    command_options: CommandOptions,
    transport_options: TransportOptions | None = None,
) -> None:
    resolved_transport = transport_options or TransportOptions()
    process = await _create_command_process(
        command_options=command_options,
        reader_limit=resolved_transport.reader_limit,
    )

    assert process.stdin is not None
    assert process.stdout is not None

    websocket_to_stdin = asyncio.create_task(
        _relay_websocket_to_stdin(websocket, process.stdin),
        name="acpremote-websocket-to-stdin",
    )
    stdout_to_websocket = asyncio.create_task(
        _relay_stdout_to_websocket(process.stdout, websocket),
        name="acpremote-stdout-to-websocket",
    )
    process_wait = asyncio.create_task(
        process.wait(),
        name="acpremote-command-wait",
    )

    tasks = {websocket_to_stdin, stdout_to_websocket, process_wait}
    try:
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for task in done:
            _raise_if_needed(task)

        if process_wait in done or stdout_to_websocket in done:
            await _close_websocket(websocket)
        if websocket_to_stdin in done and process.returncode is None:
            await _terminate_process(process, timeout=command_options.terminate_timeout)

        await _close_stdin(process.stdin)
        if pending:
            done, pending = await asyncio.wait(
                pending,
                timeout=command_options.terminate_timeout,
                return_when=asyncio.ALL_COMPLETED,
            )
            for task in done:
                _raise_if_needed(task)
    finally:
        for task in tasks:
            if task.done():
                continue
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task
        await _close_stdin(process.stdin)
        if process.returncode is None:
            await _terminate_process(process, timeout=command_options.terminate_timeout)
        await _close_websocket(websocket)


async def _create_command_process(
    *,
    command_options: CommandOptions,
    reader_limit: int,
) -> asyncio.subprocess.Process:
    stderr = None if command_options.stderr_mode == "inherit" else asyncio.subprocess.DEVNULL
    return await asyncio.create_subprocess_exec(
        *command_options.command,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=stderr,
        cwd=command_options.cwd,
        env=_build_process_env(command_options.env),
        limit=reader_limit,
    )


def _build_process_env(overrides: Mapping[str, str] | None) -> dict[str, str]:
    if overrides is None:
        return dict(os.environ)
    merged = dict(os.environ)
    merged.update(overrides)
    return merged


async def _relay_websocket_to_stdin(
    websocket: ServerConnection,
    stdin: asyncio.StreamWriter,
) -> None:
    try:
        while True:
            message = await websocket.recv()
            if isinstance(message, bytes):
                raise TypeError("binary WebSocket frames are not supported")
            stdin.write(message.encode("utf-8") + b"\n")
            await stdin.drain()
    except ConnectionClosed:
        return
    finally:
        await _close_stdin(stdin)


async def _relay_stdout_to_websocket(
    stdout: asyncio.StreamReader,
    websocket: ServerConnection,
) -> None:
    while True:
        line = await stdout.readline()
        if not line:
            return
        payload = line[:-1] if line.endswith(b"\n") else line
        await websocket.send(payload.decode("utf-8"))


async def _close_stdin(stdin: asyncio.StreamWriter) -> None:
    if stdin.is_closing():
        return
    stdin.close()
    with contextlib.suppress(BrokenPipeError, ConnectionResetError):
        await stdin.wait_closed()


async def _close_websocket(websocket: ServerConnection) -> None:
    with contextlib.suppress(ConnectionClosed):
        await websocket.close()


async def _terminate_process(process: asyncio.subprocess.Process, *, timeout: float) -> None:
    try:
        if process.returncode is not None:
            return

        with contextlib.suppress(ProcessLookupError):
            process.terminate()

        try:
            await asyncio.wait_for(process.wait(), timeout=timeout)
            return
        except TimeoutError:
            pass
        except ProcessLookupError:
            return

        if process.returncode is not None:
            return

        with contextlib.suppress(ProcessLookupError):
            process.kill()

        with contextlib.suppress(ProcessLookupError, TimeoutError):
            await asyncio.wait_for(process.wait(), timeout=timeout)
    except asyncio.CancelledError:
        if process.returncode is None:
            with contextlib.suppress(ProcessLookupError):
                process.kill()
            with contextlib.suppress(ProcessLookupError, asyncio.CancelledError, TimeoutError):
                await asyncio.wait_for(asyncio.shield(process.wait()), timeout=timeout)
        raise


def _raise_if_needed(task: asyncio.Task[object]) -> None:
    try:
        task.result()
    except asyncio.CancelledError:
        raise
    except ConnectionClosed:
        return
    except BrokenPipeError:
        return
    except ConnectionResetError:
        return
    except ProcessLookupError:
        return
    except Exception as exc:  # pragma: no cover - defensive logging path
        print(f"acpremote command relay failed: {exc}", file=sys.stderr)
        raise
