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

from spark_executor.tools.connections import (
    delete_connection,
    get_connection,
    list_connections,
    save_connection,
)
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.requests import (
    ConnectionNameRequest,
    EmptyRequest,
    GetJobLogsRequest,
    JobIdRequest,
    PendingIdRequest,
    PrepareSubmitJobRequest,
    SaveConnectionRequest,
)
from spark_executor.tools.status import get_job_status
from spark_executor.tools.submit import (
    cancel_pending_job,
    confirm_submit_job,
    get_pending_job,
    list_pending_jobs,
    prepare_submit_job,
)

app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")


@app.get("/health")
def health_check():
    return {"status": "ok"}


# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools.
# Each route takes a single Pydantic body model so tools/call (which sends args
# as JSON body) works for every tool, including those with dict-typed params
# like spark_conf.

# --- Pending submission flow (two-step submit) ---

@app.post("/prepare_submit_job")
def _prepare_submit_job(req: PrepareSubmitJobRequest):
    return prepare_submit_job(**req.model_dump())


@app.post("/confirm_submit_job")
def _confirm_submit_job(req: PendingIdRequest):
    return confirm_submit_job(pending_id=req.pending_id)


@app.post("/list_pending_jobs")
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
    return list_pending_jobs()


@app.post("/get_pending_job")
def _get_pending_job(req: PendingIdRequest):
    return get_pending_job(req.pending_id)


@app.post("/cancel_pending_job")
def _cancel_pending_job(req: PendingIdRequest):
    return cancel_pending_job(req.pending_id)


# --- Spark job tools ---

@app.post("/get_job_status")
def _get_job_status(req: JobIdRequest):
    return get_job_status(req.job_id)


@app.post("/get_job_logs")
def _get_job_logs(req: GetJobLogsRequest):
    return get_job_logs(req.job_id, tail_chars=req.tail_chars)


@app.post("/kill_job")
def _kill_job(req: JobIdRequest):
    return kill_job(req.job_id)


# --- Connection management tools ---

@app.post("/save_connection")
def _save_connection(req: SaveConnectionRequest):
    # exclude_none so we don't overwrite the function's default with explicit None
    return save_connection(**req.model_dump(exclude_none=True))


@app.post("/list_connections")
def _list_connections(_req: EmptyRequest = EmptyRequest()):
    return list_connections()


@app.post("/get_connection")
def _get_connection(req: ConnectionNameRequest):
    return get_connection(req.name)


@app.post("/delete_connection")
def _delete_connection(req: ConnectionNameRequest):
    return delete_connection(req.name)