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

Pydantic request models for the FastAPI route layer. The underlying tool
functions in tools/*.py still take keyword arguments; these models exist only
so fastapi-mcp can call the routes via tools/call (which sends args as a
JSON body) without 422-ing on dict-typed parameters like spark_conf.
"""
from pydantic import BaseModel, Field


class EmptyRequest(BaseModel):
    """Used for tools that take no arguments (list_connections, list_pending_jobs)."""
    pass


class SaveConnectionRequest(BaseModel):
    name: str
    master: str
    deploy_mode: str = "cluster"
    yarn_rm_url: str | None = None
    spark_conf: dict[str, str] | None = None


class PrepareSubmitJobRequest(BaseModel):
    connection: str
    script_path: str = Field(
        ...,
        description=(
            "Absolute path to the PySpark script inside the container's "
            "filesystem. Must point at an existing regular file. For "
            "LLM-generated code, call generate_job_file(code=...) first "
            "and pass the returned script_path here. For pre-existing "
            "files on the host, mount them via a docker volume and pass "
            "the in-container path. Returns 400 with a remediation hint "
            "if the path is missing or not a file."
        ),
    )
    queue: str = "default"
    executor_memory: str = "4G"
    executor_cores: int = 2
    num_executors: int = 2


class PendingIdRequest(BaseModel):
    pending_id: str


class JobIdRequest(BaseModel):
    job_id: str


class GetJobLogsRequest(BaseModel):
    job_id: str
    tail_chars: int = 5000


class ConnectionNameRequest(BaseModel):
    name: str


class GenerateJobFileRequest(BaseModel):
    code: str = Field(
        ...,
        description=(
            "Full PySpark source code to write to disk. Will be passed verbatim "
            "to spark-submit after the agent calls prepare_submit_job on the "
            "returned path."
        ),
    )