# coding=utf-8
"""
@Time   :2026/7/8
@Author :tao.chen

Tools for inspecting YARN applications NOT submitted through this service.
These bypass the local JobStore and require the caller to supply both the
YARN application_id and the name of a saved Connection.
"""
import json

from common.logging import logger
from spark_executor.core.yarn_client import YarnClientConfig, get_application_status, get_application_logs
from spark_executor.models import JobStatus, JobResult
from spark_executor.tools.connections import store as conn_store


def get_external_job_logs(application_id: str, connection_name: str, tail_chars: int = 5000) -> str:
    """Fetch aggregated container logs for a YARN application by ID."""
    logger.debug(f"get_external_job_logs enter application_id={application_id} connection_name={connection_name} tail_chars={tail_chars}")
    conn = conn_store.get(connection_name)
    if conn is None:
        raise KeyError(f"Connection not found: {connection_name}")
    config = YarnClientConfig.from_connection(conn)
    full = get_application_logs(application_id, config)
    tailed = full[-tail_chars:] if len(full) > tail_chars else full
    logger.info(
        f"get_external_job_logs ok application_id={application_id} connection_name={connection_name} "
        f"full_chars={len(full)} returned_chars={len(tailed)}"
    )
    return tailed


def get_external_job_status(application_id: str, connection_name: str) -> JobStatus:
    """Query YARN for an external application's current status."""
    logger.debug(f"get_external_job_status enter application_id={application_id} connection_name={connection_name}")
    conn = conn_store.get(connection_name)
    if conn is None:
        raise KeyError(f"Connection not found: {connection_name}")
    config = YarnClientConfig.from_connection(conn)
    state, raw = get_application_status(application_id, config)
    logger.info(f"get_external_job_status ok application_id={application_id} connection_name={connection_name} state={state}")
    return JobStatus(application_id=application_id, state=state, raw=raw)


def get_external_job_result(application_id: str, connection_name: str) -> JobResult:
    """Query YARN for an external application's terminal result view."""
    logger.debug(f"get_external_job_result enter application_id={application_id} connection_name={connection_name}")
    conn = conn_store.get(connection_name)
    if conn is None:
        raise KeyError(f"Connection not found: {connection_name}")
    config = YarnClientConfig.from_connection(conn)
    state, raw = get_application_status(application_id, config)
    app = json.loads(raw).get("app", {})
    result = JobResult(
        application_id=application_id,
        state=state,
        final_status=app.get("finalStatus"),
        diagnostics=app.get("diagnostics"),
        tracking_url=app.get("trackingUrl"),
        started_time=app.get("startedTime"),
        finished_time=app.get("finishedTime"),
    )
    logger.info(f"get_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}")
    return result