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/cancel/hold/release/checkpoint/restart) 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)

	// ── advanced tools (Phase 2) ──
	// These help the agent close the submit → watch → analyze → suggest loop
	// against the request/response MCP model.

	// ── scheduler_wait_job ──
	s.AddTool(mcp.NewTool("scheduler_wait_job",
		mcp.WithDescription("Poll a job until it reaches a terminal state (DONE/EXIT) or a timeout. Returns the final status and elapsed time. Use after scheduler_submit to follow a job to completion."),
		mcp.WithString("id",
			mcp.Description("Job ID to wait for"),
			mcp.Required(),
		),
		mcp.WithNumber("timeout_s", mcp.Description("Maximum seconds to wait (default 120)")),
		mcp.WithNumber("poll_interval_s", mcp.Description("Poll interval seconds (default 3)")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleWaitJob)

	// ── scheduler_analyze_failure ──
	s.AddTool(mcp.NewTool("scheduler_analyze_failure",
		mcp.WithDescription("Gather everything about a failed (EXIT) job — job info, execution hosts, and output logs — for failure analysis. Returns structured data; interpret the likely cause from the logs and job metadata."),
		mcp.WithString("id",
			mcp.Description("Job ID of the failed job"),
			mcp.Required(),
		),
		mcp.WithNumber("log_tail", mcp.Description("Max log lines to include (default 200)")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleAnalyzeFailure)

	// ── scheduler_suggest_resources ──
	s.AddTool(mcp.NewTool("scheduler_suggest_resources",
		mcp.WithDescription("Suggest resource parameters (processors, walltime, memory) for a job based on current node capacity and queue state. Returns concrete recommendations for scheduler_submit."),
		mcp.WithString("command", mcp.Description("The command the job will run (used to gauge footprint)")),
		mcp.WithString("queue", mcp.Description("Target queue (default normal)")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleSuggestResources)

	// ── scheduler_estimate_wait_time ──
	s.AddTool(mcp.NewTool("scheduler_estimate_wait_time",
		mcp.WithDescription("Estimate how long a pending job would wait before dispatch, based on current pending jobs, node capacity, and queue state. Returns a rough estimate — treat as guidance, not a guarantee."),
		mcp.WithString("queue", mcp.Description("Target queue (default normal)")),
		mcp.WithNumber("processors", mcp.Description("Number of processors the job requests (default 1)")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandleEstimateWaitTime)

	// ── scheduler_pend_reasons (bjobs -p equivalent) ──
	s.AddTool(mcp.NewTool("scheduler_pend_reasons",
		mcp.WithDescription("Explain why PEND (queued) jobs are waiting — the bjobs -p equivalent. Lists pending jobs (or one by id) and translates each job's PEND reason into a human-readable scenario (no_resource, match_fail, limit_host/global, queue_closed, host_unavail, ...)."),
		mcp.WithString("id", mcp.Description("Specific job id to explain (optional; defaults to all pending jobs)")),
		mcp.WithString("queue", mcp.Description("Filter by queue name")),
		mcp.WithString("owner", mcp.Description("Filter by job owner username")),
		mcp.WithNumber("limit", mcp.Description("Max results to return")),
		mcp.WithReadOnlyHintAnnotation(true),
		mcp.WithDestructiveHintAnnotation(false),
	), app.HandlePendReasons)

	// ── 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)

	// ── scheduler_checkpoint ──
	s.AddTool(mcp.NewTool("scheduler_checkpoint",
		mcp.WithDescription("Checkpoint a running job by ID (action 'checkpoint', CRIU). Destructive — you will be asked to confirm before execution."),
		mcp.WithString("id", mcp.Description("Job ID to checkpoint"), mcp.Required()),
		mcp.WithReadOnlyHintAnnotation(false),
		mcp.WithDestructiveHintAnnotation(true),
	), app.HandleCheckpointJob)

	// ── scheduler_restart ──
	s.AddTool(mcp.NewTool("scheduler_restart",
		mcp.WithDescription("Restart (re-queue) a finished job as a new job by ID, resuming from its last checkpoint if one exists. Returns the new job ID. Destructive — you will be asked to confirm before execution."),
		mcp.WithString("id", mcp.Description("Job ID to restart"), mcp.Required()),
		mcp.WithReadOnlyHintAnnotation(false),
		mcp.WithDestructiveHintAnnotation(true),
	), app.HandleRestartJob)

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