"""Utilities for running shell commands.

This module provides:
- `run_command`: capture stdout/stderr after process completes
- `run_command_streaming`: capture stdout/stderr while optionally streaming lines to callbacks
"""

from __future__ import annotations

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

from aida_mcp.output.policy import MAX_STREAM_LINE_BYTES
from aida_mcp.utils.artifacts import get_artifacts_root
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__)

LOG_ROOT = get_artifacts_root() / "logs"
MAX_CAPTURE_CHARS = 2_000
MAX_CAPTURE_LINES = 80


def open_streaming_log(cmd: list[str], cwd: Path) -> tuple[Path | None, Any]:
    return _open_streaming_log(cmd, cwd)


def write_stream_log_line(handle: Any, stream_name: str, line: str) -> None:
    _write_stream_log_line(handle, stream_name, line)


def finalize_stream_log(handle: Any, exit_code: int) -> None:
    _finalize_stream_log(handle, exit_code)


@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:
    """
    Run a shell command while streaming output lines to optional callbacks.

    This is useful for tools that need to parse live output for progress (e.g., Rush, Gradle tests).
    It still returns a regular `CommandResult` and writes a log file just like `run_command`.
    """
    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))

        # Consume streams concurrently, with timeout enforced on full run.
        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:
    """
    Run a shell command and return exit code, stdout, stderr.

    Args:
        cmd: Command and arguments as list
        cwd: Working directory
        timeout: Command timeout in seconds (default: 300)
        env: Optional environment variables dict to set for the command
        ctx: Optional MCP Context for progress notifications
        step_name: Optional step name for progress notifications (e.g., "Formatting", "Testing")

    Returns:
        Tuple of (exit_code, stdout, stderr)
    """
    process = None
    heartbeat_task: asyncio.Task[None] | None = None
    try:
        logger.info(f"Running command: {' '.join(cmd)} in {cwd}")

        # Send progress notification: step started
        if ctx and step_name:
            with suppress(Exception):  # Context may be closed/invalid
                await ctx.info(f"▶️  {step_name} started...")
        if tracker and tracker_label:
            with suppress(Exception):  # Context may be closed/invalid
                await tracker.heartbeat(f"{tracker_label} running", force=True)

        # Prepare environment
        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")

        logger.info(f"Command completed with exit code: {exit_code}")

        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task

        # Send progress notification: step completed
        if ctx and step_name:
            with suppress(Exception):  # Context may be closed/invalid
                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(f"Command timed out after {timeout}s")
        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 e:
        logger.error(f"Command execution failed: {e}")
        if heartbeat_task:
            heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat_task
        stderr_msg = f"Command execution failed: {e}"
        log_path = write_command_log(cmd, cwd, "", stderr_msg, -1)
        return CommandResult(-1, "", stderr_msg, log_path)


def format_command_output(
    command: str,
    exit_code: int,
    stdout: str,
    stderr: str,
) -> str:
    """
    Format command output for MCP tool response.

    Args:
        command: Command that was run
        exit_code: Exit code from command
        stdout: Standard output
        stderr: Standard error

    Returns:
        Formatted output string
    """
    result = [
        f"Command: {command}",
        f"Exit Code: {exit_code}",
        "",
    ]

    if exit_code == 0:
        result.append("✅ PASSED")
    else:
        result.append("❌ FAILED")

    if stdout.strip():
        result.extend(
            [
                "",
                "=== Output ===",
                stdout.strip(),
            ]
        )

    if stderr.strip():
        result.extend(
            [
                "",
                "=== Errors ===",
                stderr.strip(),
            ]
        )

    return "\n".join(result)


def write_command_log(
    cmd: list[str],
    cwd: Path,
    stdout: str,
    stderr: str,
    exit_code: int,
) -> Path | None:
    now = datetime.now(UTC)
    timestamp = now.strftime("%Y%m%d-%H%M%S-%f")
    command_name = Path(cmd[0]).name if cmd else "command"
    slug = _slugify(command_name)
    log_dir = LOG_ROOT / now.strftime("%Y%m%d")
    try:
        log_dir.mkdir(parents=True, exist_ok=True)
    except OSError:
        return None

    log_path = log_dir / f"{timestamp}-{slug}.log"
    header = [
        f"Timestamp: {now.isoformat()}",
        f"Command: {' '.join(cmd)}",
        f"CWD: {cwd}",
        f"Exit Code: {exit_code}",
        "",
        "--- stdout ---",
        stdout.rstrip() or "<empty>",
        "",
        "--- stderr ---",
        stderr.rstrip() or "<empty>",
        "",
    ]
    try:
        log_path.write_text("\n".join(header), encoding="utf-8")
    except OSError:
        return None
    return log_path


def _open_streaming_log(cmd: list[str], cwd: Path) -> tuple[Path | None, Any]:
    now = datetime.now(UTC)
    timestamp = now.strftime("%Y%m%d-%H%M%S-%f")
    command_name = Path(cmd[0]).name if cmd else "command"
    slug = _slugify(command_name)
    log_dir = LOG_ROOT / now.strftime("%Y%m%d")
    try:
        log_dir.mkdir(parents=True, exist_ok=True)
    except OSError:
        return None, None

    log_path = log_dir / f"{timestamp}-{slug}.log"
    try:
        handle = log_path.open("w", encoding="utf-8")
        handle.write(f"Timestamp: {now.isoformat()}\n")
        handle.write(f"Command: {' '.join(cmd)}\n")
        handle.write(f"CWD: {cwd}\n")
        handle.write("Exit Code: <pending>\n\n")
        return log_path, handle
    except OSError:
        return None, None


def _write_stream_log_line(handle: Any, stream_name: str, line: str) -> None:
    if handle is None:
        return
    try:
        handle.write(f"[{stream_name}] {line}\n")
    except OSError:
        return


def _finalize_stream_log(handle: Any, exit_code: int) -> None:
    if handle is None:
        return
    try:
        handle.write(f"\nExit Code: {exit_code}\n")
        handle.flush()
        handle.close()
    except OSError:
        return


def _slugify(value: str) -> str:
    normalized = re.sub(r"[^A-Za-z0-9._-]+", "-", value).strip("-")
    return normalized or "command"


async def _heartbeat_loop(tracker: ProgressTracker, label: str) -> None:
    """Emit tracker heartbeats while the process runs."""

    try:
        while True:
            await asyncio.sleep(tracker.heartbeat_interval)
            try:
                await tracker.heartbeat(f"{label} running")
            except Exception:
                # Context may be closed/invalid - stop heartbeating
                return
    except asyncio.CancelledError:
        return
