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

from common.logging import logger


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]":
    logger.debug(f"yarn _run exec: {cmd}")
    result = subprocess.run(cmd, capture_output=True, text=True, errors="replace")
    logger.debug(
        f"yarn _run done rc={result.returncode} "
        f"stdout_len={len(result.stdout)} stderr_len={len(result.stderr)}"
    )
    return result


def get_application_status(application_id: str) -> tuple[str, str]:
    proc = _run(["yarn", "application", "-status", application_id])
    if proc.returncode != 0:
        logger.error(
            f"yarn application -status failed (rc={proc.returncode}) for "
            f"{application_id}: {proc.stderr[:500]}"
        )
        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}")
    state = match.group(1)
    logger.info(f"yarn status {application_id} -> {state}")
    return state, proc.stdout


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


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