# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import YarnClientConfig, get_application_logs
from spark_executor.tools.connections import store as conn_store

store = JobStore()


def _unknown_job_error(uid: str, external_tool_hint: str | None = None) -> Exception:
    """Build the right error for a not-found job.

    The agent gets this from confirm_submit_job's response:
        {"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
    and routinely confuses which to pass here. We differentiate two cases:

    - `uid` looks like a YARN application_id (starts with 'application_')
      AND the caller passed an `external_tool_hint`: the YARN app likely
      exists, this tool just can't serve it because it was not submitted
      through this service. Raise ValueError (-> 400 via the FastAPI
      handler) pointing the agent at the right external tool.
    - Otherwise: no local record of either form of id. Raise KeyError
      (-> 404). Spelling out what both IDs look like saves a round trip.
    """
    if uid.startswith("application_") and external_tool_hint:
        return ValueError(
            f"job_id={uid!r} looks like a YARN application_id (starts with "
            f"'application_'), but this tool only works for jobs submitted "
            f"through this MCP service (no local JobStore record). For YARN "
            f"applications not submitted here, use {external_tool_hint} "
            f"instead. (If you actually submitted this job through this "
            f"service, pass the local job_id — it is a 12-char hex like "
            f"'a1b2c3d4e5f6'.)"
        )
    return KeyError(
        f"No Job found for id={uid!r} (neither as job_id nor as "
        f"application_id). Pass the job_id from confirm_submit_job's "
        f"response — it is a 12-char hex like 'a1b2c3d4e5f6'. The "
        f"application_id is the YARN ID, e.g. 'application_17400000001_0001'."
    )


def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
    """Fetch aggregated container logs for a Spark job.

    `job_id` accepts either the local job_id (returned by
    confirm_submit_job) or the YARN application_id — both are looked up
    against the same Job record.
    """
    logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}")
    job = store.get_either(job_id)
    if job is None:
        raise _unknown_job_error(
            job_id,
            external_tool_hint="get_external_job_logs(application_id, connection_name, tail_chars)",
        )
    conn = conn_store.get(job.connection)
    if conn is None:
        raise KeyError(f"Connection not found: {job.connection}")
    config = YarnClientConfig.from_connection(conn)
    full = get_application_logs(job.application_id, config)
    tailed = full[-tail_chars:] if len(full) > tail_chars else full
    logger.info(
        f"get_job_logs ok job_id={job.job_id} application_id={job.application_id} "
        f"full_chars={len(full)} returned_chars={len(tailed)}"
    )
    return tailed