# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import re
import subprocess


class YarnError(Exception):
    """Raised when a yarn CLI invocation fails."""


_STATE_RE = re.compile(r"State\s*:\s*(\S+)")


def _run(cmd: list[str]) -> "subprocess.CompletedProcess[str]":
    return subprocess.run(cmd, capture_output=True, text=True, errors="replace")


def get_application_status(application_id: str) -> tuple[str, str]:
    proc = _run(["yarn", "application", "-status", application_id])
    if proc.returncode != 0:
        raise YarnError(
            f"yarn application -status failed (rc={proc.returncode}): {proc.stderr}"
        )
    match = _STATE_RE.search(proc.stdout)
    if not match:
        raise YarnError(f"Could not parse YARN state from output: {proc.stdout!r}")
    return match.group(1), proc.stdout


def get_application_logs(application_id: str) -> str:
    proc = _run(["yarn", "logs", "-applicationId", application_id])
    if proc.returncode != 0:
        raise YarnError(
            f"yarn logs failed (rc={proc.returncode}): {proc.stderr}"
        )
    return proc.stdout


def kill_application(application_id: str) -> None:
    proc = _run(["yarn", "application", "-kill", application_id])
    if proc.returncode != 0:
        raise YarnError(
            f"yarn application -kill failed (rc={proc.returncode}): {proc.stderr}"
        )