# Copyright (C) 2026 GoodData Corporation

"""Command execution helpers (streaming and non-streaming)."""

from __future__ import annotations

import asyncio
import logging
import os
from collections import deque
from collections.abc import Awaitable, Callable
from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

from aida_mcp.output.policy import MAX_STREAM_LINE_BYTES
from aida_mcp.utils.command_runner_logs import (
    finalize_stream_log,
    open_streaming_log,
    write_command_log,
    write_stream_log_line,
)
from aida_mcp.utils.stream_reader import read_stream_chunked

if TYPE_CHECKING:
    from mcp.server.fastmcp import Context

    from aida_mcp.utils.progress import ProgressTracker

logger = logging.getLogger(__name__)

MAX_CAPTURE_CHARS = 2_000
MAX_CAPTURE_LINES = 80


@dataclass(slots=True)
class CommandResult:
    exit_code: int
    stdout: str
    stderr: str
    log_path: Path | None


async def run_command_streaming(
    cmd: list[str],
    cwd: Path,
    *,
    timeout: int = 300,
    env: dict[str, str] | None = None,
    on_stdout_line: Callable[[str], Awaitable[None]] | None = None,
    on_stderr_line: Callable[[str], Awaitable[None]] | None = None,
    ctx: Context[Any, Any, Any] | None = None,
    step_name: str | None = None,
    tracker: ProgressTracker | None = None,
    tracker_label: str | None = None,
) -> CommandResult:
    process: asyncio.subprocess.Process | None = None
    heartbeat_task: asyncio.Task[None] | None = None
    stdout_lines: deque[str] = deque()
    stderr_lines: deque[str] = deque()
    stdout_chars = 0
    stderr_chars = 0
    log_path, log_handle = open_streaming_log(cmd, cwd)

    async def _read_stream(
        stream: asyncio.StreamReader | None,
        *,
        sink: deque[str],
        on_line: Callable[[str], Awaitable[None]] | None,
        stream_name: str,
    ) -> None:
        nonlocal stdout_chars, stderr_chars
        if sink is stdout_lines:
            stdout_chars = await read_stream_chunked(
                stream,
                stream_name=stream_name,
                log_handle=log_handle,
                write_log_line=write_stream_log_line,
                on_line=on_line,
                max_capture_lines=MAX_CAPTURE_LINES,
                max_capture_chars=MAX_CAPTURE_CHARS,
                max_line_bytes=MAX_STREAM_LINE_BYTES,
                sink=sink,
                chars_counter=stdout_chars,
            )
        else:
            stderr_chars = await read_stream_chunked(
                stream,
                stream_name=stream_name,
                log_handle=log_handle,
                write_log_line=write_stream_log_line,
                on_line=on_line,
                max_capture_lines=MAX_CAPTURE_LINES,
                max_capture_chars=MAX_CAPTURE_CHARS,
                max_line_bytes=MAX_STREAM_LINE_BYTES,
                sink=sink,
                chars_counter=stderr_chars,
            )

    try:
        logger.info("Running (streaming) command: %s in %s", " ".join(cmd), cwd)
        if ctx and step_name:
            with suppress(Exception):
                await ctx.info(f"▶️  {step_name} started...")
        if tracker and tracker_label:
            with suppress(Exception):
                await tracker.heartbeat(f"{tracker_label} running", force=True)
        process_env = None
        if env:
            process_env = os.environ.copy()
            process_env.update(env)
        process = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=str(cwd),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=process_env,
        )
        if tracker and tracker_label:
            heartbeat_task = asyncio.create_task(_heartbeat_loop(tracker, tracker_label))
        await asyncio.wait_for(
            asyncio.gather(
                _read_stream(process.stdout, sink=stdout_lines, on_line=on_stdout_line, stream_name="stdout"),
                _read_stream(process.stderr, sink=stderr_lines, on_line=on_stderr_line, stream_name="stderr"),
                process.wait(),
            ),
            timeout=timeout,
        )
        exit_code = process.returncode or 0
        stdout_str = "\n".join(stdout_lines)
        stderr_str = "\n".join(stderr_lines)
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        if ctx and step_name:
            with suppress(Exception):
                if exit_code == 0:
                    await ctx.info(f"✅ {step_name} completed")
                else:
                    await ctx.info(f"❌ {step_name} failed (exit code: {exit_code})")
        finalize_stream_log(log_handle, exit_code)
        return CommandResult(exit_code, stdout_str, stderr_str, log_path)
    except TimeoutError:
        logger.error("Streaming command timed out after %ss", timeout)
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        if process:
            process.kill()
            await process.wait()
        stdout_str = "\n".join(stdout_lines)
        stderr_msg = f"Command timed out after {timeout} seconds"
        finalize_stream_log(log_handle, -1)
        return CommandResult(-1, stdout_str, stderr_msg, log_path)
    except Exception as exc:
        logger.error("Streaming command execution failed: %s", exc)
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        stdout_str = "\n".join(stdout_lines)
        stderr_msg = f"Command execution failed: {exc}"
        finalize_stream_log(log_handle, -1)
        return CommandResult(-1, stdout_str, stderr_msg, log_path)


async def run_command(
    cmd: list[str],
    cwd: Path,
    timeout: int = 300,
    env: dict[str, str] | None = None,
    ctx: Context[Any, Any, Any] | None = None,
    step_name: str | None = None,
    tracker: ProgressTracker | None = None,
    tracker_label: str | None = None,
) -> CommandResult:
    process = None
    heartbeat_task: asyncio.Task[None] | None = None
    try:
        logger.info("Running command: %s in %s", " ".join(cmd), cwd)
        if ctx and step_name:
            with suppress(Exception):
                await ctx.info(f"▶️  {step_name} started...")
        if tracker and tracker_label:
            with suppress(Exception):
                await tracker.heartbeat(f"{tracker_label} running", force=True)
        process_env = None
        if env:
            process_env = os.environ.copy()
            process_env.update(env)
        process = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=str(cwd),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=process_env,
        )
        if tracker and tracker_label:
            heartbeat_task = asyncio.create_task(_heartbeat_loop(tracker, tracker_label))
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
        exit_code = process.returncode or 0
        stdout_str = stdout.decode("utf-8", errors="replace")
        stderr_str = stderr.decode("utf-8", errors="replace")
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        if ctx and step_name:
            with suppress(Exception):
                if exit_code == 0:
                    await ctx.info(f"✅ {step_name} completed")
                else:
                    await ctx.info(f"❌ {step_name} failed (exit code: {exit_code})")
        log_path = write_command_log(cmd, cwd, stdout_str, stderr_str, exit_code)
        return CommandResult(exit_code, stdout_str, stderr_str, log_path)
    except TimeoutError:
        logger.error("Command timed out after %ss", timeout)
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        if process:
            process.kill()
            await process.wait()
        stderr_msg = f"Command timed out after {timeout} seconds"
        log_path = write_command_log(cmd, cwd, "", stderr_msg, -1)
        return CommandResult(-1, "", stderr_msg, log_path)
    except Exception as exc:
        logger.error("Command execution failed: %s", exc)
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        stderr_msg = f"Command execution failed: {exc}"
        log_path = write_command_log(cmd, cwd, "", stderr_msg, -1)
        return CommandResult(-1, "", stderr_msg, log_path)


async def _heartbeat_loop(tracker: ProgressTracker, label: str) -> None:
    try:
        while True:
            await asyncio.sleep(tracker.heartbeat_interval)
            try:
                await tracker.heartbeat(f"{label} running")
            except Exception:
                return
    except asyncio.CancelledError:
        return
