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

The MCP-exposed `write_job_file` tool. Despite the historical name
`generate_job_file`, this tool does NOT generate PySpark code — the
calling LLM writes the code, and this tool persists it to disk so
`spark-submit` can see it. The rename (2026-06-29) tightens the naming
to match what the tool actually does, eliminating the "I thought this
would generate code for me" confusion.

It also runs the SQL safety policy (SELECT/INSERT only) on the code
BEFORE writing, so the agent gets immediate feedback if it slipped a
DROP / DELETE / UPDATE past its own code-generation step.
"""
from common.logging import logger
from common.sql_guard import validate_pyspark_code
# Internal helper already named write_job_file (it just writes bytes to
# disk; no SQL guard). Import with an alias to avoid the name collision
# with the MCP-exposed function below.
from spark_executor.core.job_writer import write_job_file as _write_to_disk


class SqlGuardViolation(ValueError):
    """Raised when write_job_file receives code that violates the SQL
    safety policy. -> HTTP 400 via the FastAPI ValueError handler."""
    pass


def write_job_file(code: str) -> dict[str, str]:
    """Persist a PySpark code string to disk and return its absolute path.

    Validates the code against the SQL safety policy (SELECT/INSERT only)
    BEFORE writing, so the agent gets immediate feedback rather than
    learning at prepare_submit_job time. Use the returned path as the
    `script_path` argument of prepare_submit_job.

    The output directory is controlled by the SPARK_EXECUTOR_JOBS_DIR env var
    (default: ./data/jobs/). Each call writes to a fresh timestamped file
    under that directory, so repeated calls never overwrite one another.
    """
    logger.debug(f"write_job_file enter code_bytes={len(code)}")
    offenses = validate_pyspark_code(code)
    if offenses:
        logger.warning(
            f"write_job_file rejected: SQL policy violation(s): {offenses}"
        )
        raise SqlGuardViolation(
            f"PySpark code violates SQL safety policy. "
            f"Only SELECT and INSERT statements are allowed. "
            f"Found forbidden statement(s): {offenses}. "
            f"Rewrite the code to use only SELECT/INSERT."
        )
    path = _write_to_disk(code)
    logger.info(f"write_job_file ok script_path={path}")
    return {"script_path": path}