# 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.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.

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

@app.post("/prepare_submit_job")
def _prepare_submit_job(connection: str, script_path: str, queue: str = "default",
                        executor_memory: str = "4G", executor_cores: int = 2,
                        num_executors: int = 2):
    return prepare_submit_job(
        connection=connection,
        script_path=script_path,
        queue=queue,
        executor_memory=executor_memory,
        executor_cores=executor_cores,
        num_executors=num_executors,
    )


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


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


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


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


# --- Spark job tools ---

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


@app.post("/get_job_logs")
def _get_job_logs(job_id: str, tail_chars: int = 5000):
    return get_job_logs(job_id, tail_chars=tail_chars)


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


# --- Connection management tools ---

@app.post("/save_connection")
def _save_connection(name: str, master: str, deploy_mode: str = "cluster",
                     yarn_rm_url: str | None = None,
                     spark_conf: dict[str, str] | None = None):
    return save_connection(
        name=name,
        master=master,
        deploy_mode=deploy_mode,
        yarn_rm_url=yarn_rm_url,
        spark_conf=spark_conf,
    )


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


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


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