# Copyright 2026 Datastrato, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Stdio transport for ADP clients.

Starts an ADP server as a subprocess and communicates
via stdin/stdout using newline-delimited JSON.
"""

import asyncio
import logging
import os
import signal
import sys

from adp_sdk.transports.base import Transport

logger = logging.getLogger(__name__)

# Grace period before killing the subprocess
_SHUTDOWN_TIMEOUT_S = 5.0


class StdioClientTransport(Transport):
    """Stdio transport that launches an ADP server subprocess.

    Communicates via newline-delimited JSON over the subprocess's
    stdin (write) and stdout (read). Stderr output from the subprocess
    is forwarded to the Python logging system.

    Usage::

        transport = StdioClientTransport(
            "python",
            ["-m", "adp_hypervisor", "--config", "/path/to/manifests"],
        )
        async with transport:
            await transport.write_message('{"jsonrpc":"2.0",...}')
            response = await transport.read_message()
    """

    def __init__(
        self,
        command: str,
        args: list[str] | None = None,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
    ) -> None:
        self._command = command
        self._args = args or []
        self._env = env
        self._cwd = cwd
        self._process: asyncio.subprocess.Process | None = None
        self._stderr_task: asyncio.Task[None] | None = None

    async def start(self) -> None:
        """Start the subprocess and begin reading stderr.

        Raises:
            RuntimeError: If the transport is already started.
        """
        if self._process is not None:
            raise RuntimeError("Transport already started")

        # Merge custom env with current environment
        merged_env: dict[str, str] | None = None
        if self._env is not None:
            merged_env = {**os.environ, **self._env}

        self._process = await asyncio.create_subprocess_exec(
            self._command,
            *self._args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=merged_env,
            cwd=self._cwd,
        )

        self._stderr_task = asyncio.create_task(
            self._read_stderr(),
            name="stdio-client-stderr",
        )

        logger.info(
            "Started subprocess: %s %s (pid=%d)",
            self._command,
            " ".join(self._args),
            self._process.pid,
        )

    async def read_message(self) -> str | None:
        """Read a newline-delimited JSON message from subprocess stdout.

        Returns:
            The message string (without trailing newline), or None if EOF.

        Raises:
            RuntimeError: If the transport is not started.
        """
        if self._process is None or self._process.stdout is None:
            raise RuntimeError("Transport not started; call start() first")

        while True:
            line = await self._process.stdout.readline()
            if not line:
                logger.debug("Subprocess stdout reached EOF")
                return None

            message = line.decode("utf-8").strip()
            if not message:
                continue

            logger.debug("Received message: %s", message[:200])
            return message

    async def write_message(self, message: str) -> None:
        """Write a newline-delimited JSON message to subprocess stdin.

        Args:
            message: The JSON string to send.

        Raises:
            RuntimeError: If the transport is not started.
        """
        if self._process is None or self._process.stdin is None:
            raise RuntimeError("Transport not started; call start() first")

        data = message.rstrip("\n") + "\n"
        self._process.stdin.write(data.encode("utf-8"))
        await self._process.stdin.drain()
        logger.debug("Sent message: %s", message[:200])

    async def close(self) -> None:
        """Close the transport and terminate the subprocess."""
        if self._process is None:
            return

        process = self._process

        # Cancel stderr reader first (while self._process is still set)
        if self._stderr_task is not None:
            self._stderr_task.cancel()
            try:
                await self._stderr_task
            except asyncio.CancelledError:
                pass
            self._stderr_task = None

        self._process = None

        # Close stdin to signal the subprocess to exit
        if process.stdin is not None:
            try:
                process.stdin.close()
                await process.stdin.wait_closed()
            except Exception:
                logger.debug("Error closing subprocess stdin", exc_info=True)

        # Wait for graceful exit, then force-kill
        if process.returncode is None:
            try:
                sig = signal.SIGTERM if sys.platform != "win32" else signal.SIGINT
                process.send_signal(sig)
                await asyncio.wait_for(process.wait(), timeout=_SHUTDOWN_TIMEOUT_S)
            except TimeoutError:
                logger.warning("Subprocess did not exit gracefully, killing (pid=%d)", process.pid)
                process.kill()
                await process.wait()
            except ProcessLookupError:
                pass

        logger.info("Subprocess terminated (pid=%d, returncode=%s)", process.pid, process.returncode)

    async def _read_stderr(self) -> None:
        """Forward subprocess stderr lines to logging."""
        assert self._process is not None
        assert self._process.stderr is not None

        try:
            while True:
                line = await self._process.stderr.readline()
                if not line:
                    break
                text = line.decode("utf-8", errors="replace").rstrip()
                if text:
                    logger.info("[server stderr] %s", text)
        except asyncio.CancelledError:
            pass
        except Exception:
            logger.debug("Error reading subprocess stderr", exc_info=True)
