package main

import (
	"log"
	"os"

	mcps "github.com/mark3labs/mcp-go/server"
	"github.com/mark3labs/mcp-go/mcp"

	"mcp-server/internal"
)

// serverInstructions are returned to MCP clients (LLMs like Claude Code) in
// the initialize handshake and are the standard way to steer how the model
// uses the server. The key guidance: the status tools return operational
// snapshots, not incident triggers — a node reporting UNAVAIL or a job in
// PEND/EXIT/UNKWN is normal state. Without this, the model treats every
// non-"OK" value as a problem and floods the user with unsolicited
// diagnostic advice when they only asked to view status.
const serverInstructions = "XCE scheduler agent — cluster status, not an incident console.\n\n" +
	"Status tools (scheduler_list_nodes, scheduler_list_jobs, scheduler_get_job, scheduler_cluster_metrics) " +
	"return informational snapshots. A node showing UNAVAIL/offline, or a job in PEND/EXIT/UNKWN/TRANSIENT, is " +
	"normal operational state, not an error condition.\n\n" +
	"When the user asks to view or check status, report the data factually and concisely. Do NOT speculate about " +
	"root causes, do NOT reference \"known bugs\", and do NOT propose diagnostic or repair actions (checking logs, " +
	"restarting daemons, opening tickets, ssh'ing to nodes, etc.) unless the user explicitly asks for help or " +
	"remediation. If you believe action may be warranted, say so in at most one line and let the user decide.\n\n" +
	"Write tools (scheduler_submit/hold/release/cancel) are destructive and require the user's explicit intent."

func main() {
	baseURL := os.Getenv("XCE_API_URL")
	token := os.Getenv("XCE_API_KEY")

	if baseURL == "" {
		baseURL = "https://localhost:8080"
		log.Printf("XCE_API_URL not set, defaulting to %s", baseURL)
	}
	if token == "" {
		log.Fatal("XCE_API_KEY environment variable is required")
	}

	app, err := internal.NewApp(baseURL, token)
	if err != nil {
		log.Fatalf("failed to initialize: %v", err)
	}

	s := mcps.NewMCPServer(
		"scheduler-agent",
		"0.2.0",
		mcps.WithResourceCapabilities(true, true),
		mcps.WithLogging(),
		mcps.WithInstructions(serverInstructions),
	)

	// Read-only tools are annotated as such so the MCP client does not prompt
	// for confirmation on them (NewTool defaults to destructive=true).

	// ── scheduler_list_jobs ──
	s.AddTool(mcp.NewTool("scheduler_list_jobs",
		mcp.WithDescription("List jobs in the cluster. Supports filtering by queue, owner, status, and pagination. Status snapshot: report results as-is; job states (PEND/RUN/EXIT/DONE/etc.) are informational, not errors to investigate unprompted."),
		mcp.WithString("queue", mcp.Description("Filter by queue name")),
		mcp.WithString("owner", mcp.Description("Filter by job owner username")),
		mcp.WithString("status", mcp.Description("Filter by status: pending|running|succeeded|failed|cancelled|suspended|unknown")),
		mcp.WithNumber("limit", mcp.Description("Max results to return")),
		mcp.WithString("page_token", mcp.Description("Pagination token from a prior response")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleListJobs)

	// ── scheduler_get_job ──
	s.AddTool(mcp.NewTool("scheduler_get_job",
		mcp.WithDescription("Get detailed information about a specific job by its ID."),
		mcp.WithString("id",
			mcp.Description("Job ID"),
			mcp.Required(),
		),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleGetJob)

	// ── scheduler_list_queues ──
	s.AddTool(mcp.NewTool("scheduler_list_queues",
		mcp.WithDescription("List all queues/partitions and their current state (active, inactive, draining)."),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleListQueues)

	// ── scheduler_list_nodes ──
	s.AddTool(mcp.NewTool("scheduler_list_nodes",
		mcp.WithDescription("List all compute nodes with their resource usage and status. Status snapshot: report node states (OK/UNAVAIL/etc.) as-is — a non-OK node is informational, not an error; do not offer diagnostic/repair suggestions unless the user asks."),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleListNodes)

	// ── scheduler_job_logs ──
	s.AddTool(mcp.NewTool("scheduler_job_logs",
		mcp.WithDescription("Retrieve the output logs of a job."),
		mcp.WithString("id",
			mcp.Description("Job ID"),
			mcp.Required(),
		),
		mcp.WithNumber("tail", mcp.Description("Number of lines from the end to return (0 = all)")),
		mcp.WithNumber("offset_bytes", mcp.Description("Byte offset to start reading from")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleJobLogs)

	// ── scheduler_cluster_metrics ──
	s.AddTool(mcp.NewTool("scheduler_cluster_metrics",
		mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization. Status snapshot: report as-is; do not speculate about causes or propose actions unless the user asks."),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleClusterMetrics)

	// ── write tools (Phase 1) ──
	// Write tools are annotated destructive so the MCP client (e.g. Claude
	// Code) prompts the user for confirmation before executing them — the
	// design doc's "写操作默认 require confirm" at the layer MCP supports it.

	// ── scheduler_submit ──
	s.AddTool(mcp.NewTool("scheduler_submit",
		mcp.WithDescription("Submit a job. REVIEW REQUIRED: call with dry_run=true first to inspect the resolved job spec, then re-call with dry_run=false to actually submit."),
		mcp.WithString("template", mcp.Description("Job template (default 'general_serial' 通用串行脚本)."),
			mcp.Enum(internal.TemplateGeneralSerial), mcp.DefaultString(internal.TemplateGeneralSerial)),
		mcp.WithBoolean("dry_run", mcp.Description("Review mode: validate and return the resolved job spec without submitting."), mcp.DefaultBool(false)),
		mcp.WithString("command", mcp.Description("Shell command to run (required)."), mcp.Required()),
		mcp.WithString("name", mcp.Description("Job name; array syntax 'myjob[1-10]' supported.")),
		mcp.WithString("queue", mcp.Description("Queue (default 'normal').")),
		mcp.WithString("project", mcp.Description("Project name.")),
		mcp.WithNumber("cores", mcp.Description("Number of processors.")),
		mcp.WithString("walltime", mcp.Description("Walltime limit 'HH:MM', e.g. '02:30'.")),
		mcp.WithString("memory", mcp.Description("Memory limit, e.g. '4 GB', '2048 MB'.")),
		mcp.WithString("cwd", mcp.Description("Working directory.")),
		mcp.WithString("env", mcp.Description("Comma-separated KEY=VALUE environment variables.")),
		mcp.WithString("out_file", mcp.Description("Standard output file.")),
		mcp.WithString("err_file", mcp.Description("Standard error file.")),
		mcp.WithString("description", mcp.Description("Job description.")),
		mcp.WithReadOnlyHintAnnotation(false),
		mcp.WithDestructiveHintAnnotation(true),
	), app.HandleSubmitJob)

	// ── scheduler_cancel ──
	s.AddTool(mcp.NewTool("scheduler_cancel",
		mcp.WithDescription("Cancel (kill) a job by ID. Destructive — you will be asked to confirm before execution."),
		mcp.WithString("id", mcp.Description("Job ID to cancel"), mcp.Required()),
		mcp.WithReadOnlyHintAnnotation(false),
		mcp.WithDestructiveHintAnnotation(true),
	), app.HandleCancelJob)

	// ── scheduler_hold ──
	s.AddTool(mcp.NewTool("scheduler_hold",
		mcp.WithDescription("Hold (suspend/stop) a job by ID. Destructive — you will be asked to confirm before execution."),
		mcp.WithString("id", mcp.Description("Job ID to hold"), mcp.Required()),
		mcp.WithReadOnlyHintAnnotation(false),
		mcp.WithDestructiveHintAnnotation(true),
	), app.HandleHoldJob)

	// ── scheduler_release ──
	s.AddTool(mcp.NewTool("scheduler_release",
		mcp.WithDescription("Release (resume) a held job by ID. Destructive — you will be asked to confirm before execution."),
		mcp.WithString("id", mcp.Description("Job ID to release"), mcp.Required()),
		mcp.WithReadOnlyHintAnnotation(false),
		mcp.WithDestructiveHintAnnotation(true),
	), app.HandleReleaseJob)

	log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL)
	if err := mcps.ServeStdio(s); err != nil {
		log.Fatalf("server error: %v", err)
	}
}