package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

// Command line helpers used by the installer and by people on the machine:
//
//	wxch-agent pull <runtime> <model> [--store auto|internal|external]
//	                                      download a model through a runtime (progress on stdout); auto, the default,
//	                                      puts it on the external disk when one is connected, except the starter models
//	wxch-agent status [-json]             which runtimes are ready, which models are installed
//	wxch-agent selftest [runtime[=model]] load a model in each runtime and ask it one question
//	wxch-agent modules list|install|remove  the optional AI modules (image, video, speech); see cli_modules.go
//
// They run as the same user as the agent (the installer uses sudo -u), so file ownership and the
// vLLM runner's peer-uid check behave exactly as in normal operation.

func isCLICommand(s string) bool {
	return s == "pull" || s == "status" || s == "selftest" || s == "modules"
}

func runCLI(cmd string, args []string) int {
	switch cmd {
	case "pull":
		return cliPull(args)
	case "status":
		return cliStatus(args)
	case "selftest":
		return cliSelftest(args)
	case "modules":
		return cliModules(args)
	}
	return 2
}

// serverBase is the https address of the control server this machine belongs to ("" if unknown).
func serverBase() string {
	ws := os.Getenv("WXCH_SERVER")
	if ws == "" {
		ws = loadNodeConfig().Server
	}
	ws = strings.TrimSuffix(strings.TrimSuffix(ws, "/ws/agent"), "/")
	ws = strings.Replace(strings.Replace(ws, "wss://", "https://", 1), "ws://", "http://", 1)
	return ws
}

// installHint tells the user how to (re)run the installer for this server.
func installHint() string {
	b := serverBase()
	if b == "" {
		b = "https://<your-control-host>"
	}
	return "re-run the installer: curl -fsSL " + b + "/get | sh"
}

func cliPull(args []string) int {
	store := "auto"
	var pos []string
	for i := 0; i < len(args); i++ {
		if args[i] == "--store" && i+1 < len(args) {
			store = args[i+1]
			i++
		} else {
			pos = append(pos, args[i])
		}
	}
	args = pos
	if len(args) != 2 || (store != "auto" && store != "internal" && store != "external") {
		fmt.Fprintln(os.Stderr, "usage: wxch-agent pull <ollama|lmstudio|vllm> <model> [--store auto|internal|external]")
		return 2
	}
	rt := runtimeByID(args[0])
	if rt == nil {
		fmt.Fprintf(os.Stderr, "unknown runtime %q\n", args[0])
		return 2
	}
	// vLLM only needs files on disk to download, the others need their server running
	if rt.ID() != "vllm" {
		if err := rt.Ready(); err != nil {
			fmt.Fprintf(os.Stderr, "%s is not ready: %v\n", rt.ID(), err)
			return 1
		}
	}
	var last string
	var lastAt time.Time
	err := rt.Pull(context.Background(), args[1], func(status string, total, done int64) {
		line := status
		if total > 0 {
			line = fmt.Sprintf("%s  %.0f%%  %.2f/%.2f GB", status, 100*float64(done)/float64(total), float64(done)/1e9, float64(total)/1e9)
		}
		if line != last && time.Since(lastAt) > 2*time.Second {
			fmt.Println(line)
			last, lastAt = line, time.Now()
		}
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, "download failed:", err)
		return 1
	}
	fmt.Printf("downloaded %s via %s\n", args[1], rt.ID())
	// the long-running agent watches the disk; this command has to look for itself
	cfg := loadNodeConfig()
	setExternalDisks(probeAll(cfg))
	if store == "auto" && isMainModel(cfg, args[1]) {
		store = "internal" // nobody can answer the 20 s question here: the installer asks before it calls this
	}
	if err := placeAfterPull(rt.ID(), args[1], store, func(status string, total, done int64) {
		if total > 0 {
			fmt.Printf("%s  %.0f%%\n", status, 100*float64(done)/float64(total))
		} else {
			fmt.Println(status)
		}
	}); err != nil {
		fmt.Fprintln(os.Stderr, "warning: could not place the model on the external disk:", err)
		if store == "external" {
			return 1
		}
	}
	return 0
}

type runtimeReport struct {
	Runtime   string           `json:"runtime"`
	Ready     bool             `json:"ready"`
	Error     string           `json:"error,omitempty"`
	Installed map[string]int64 `json:"installed,omitempty"`
	Loaded    []string         `json:"loaded,omitempty"`
}

func collectStatus() []runtimeReport {
	var out []runtimeReport
	for _, rt := range allRuntimes {
		r := runtimeReport{Runtime: rt.ID()}
		if err := rt.Ready(); err != nil {
			r.Error = err.Error()
		} else {
			r.Ready = true
			r.Installed, _ = rt.Installed()
			if l, err := rt.Loaded(); err == nil {
				for n := range l {
					r.Loaded = append(r.Loaded, n)
				}
				sort.Strings(r.Loaded)
			}
		}
		out = append(out, r)
	}
	return out
}

func cliStatus(args []string) int {
	fs := flag.NewFlagSet("status", flag.ExitOnError)
	asJSON := fs.Bool("json", false, "machine-readable output")
	fs.Parse(args)
	rep := collectStatus()
	allReady := true
	if *asJSON {
		json.NewEncoder(os.Stdout).Encode(rep)
	}
	for _, r := range rep {
		if !r.Ready {
			allReady = false
		}
		if *asJSON {
			continue
		}
		if r.Ready {
			names := make([]string, 0, len(r.Installed))
			for n, sz := range r.Installed {
				names = append(names, fmt.Sprintf("%s (%.2f GB)", n, float64(sz)/1e9))
			}
			sort.Strings(names)
			fmt.Printf("%-9s READY    models: %s\n", r.Runtime, strings.Join(names, ", "))
		} else {
			fmt.Printf("%-9s NOT READY  %s\n", r.Runtime, r.Error)
		}
	}
	if !allReady {
		return 1
	}
	return 0
}

// smallestInstalled returns the installed model with the smallest size ("" if none).
func smallestInstalled(inst map[string]int64) string {
	best, bestSize := "", int64(-1)
	for n, sz := range inst {
		if bestSize < 0 || sz < bestSize || (sz == bestSize && n < best) {
			best, bestSize = n, sz
		}
	}
	return best
}

// askOnce sends one tiny chat request to a model through its runtime and returns the reply text.
func askOnce(ctx context.Context, base, model string) (string, error) {
	body, _ := json.Marshal(map[string]any{
		"model": model, "temperature": 0, "max_tokens": 64,
		"messages": []map[string]string{{"role": "user", "content": "Reply with the single word: OK"}},
	})
	req, _ := http.NewRequestWithContext(ctx, "POST", base+"/v1/chat/completions", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw[:min(len(raw), 200)])))
	}
	var out struct {
		Choices []struct {
			Message struct {
				Content          string `json:"content"`
				ReasoningContent string `json:"reasoning_content"`
				Reasoning        string `json:"reasoning"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(raw, &out); err != nil || len(out.Choices) == 0 {
		return "", fmt.Errorf("unexpected answer: %s", strings.TrimSpace(string(raw[:min(len(raw), 200)])))
	}
	m := out.Choices[0].Message
	text := strings.TrimSpace(m.Content)
	if text == "" {
		text = strings.TrimSpace(m.ReasoningContent + m.Reasoning) // thinking models may only think within the token limit
	}
	if text == "" {
		return "", fmt.Errorf("the model returned an empty answer")
	}
	return text, nil
}

func oneLine(s string, n int) string {
	s = strings.Join(strings.Fields(s), " ")
	if len(s) > n {
		s = s[:n] + "…"
	}
	return s
}

// cliSelftest loads one model per runtime, asks it a question, and unloads it again.
// Arguments are runtime or runtime=model; with no arguments every ready runtime is tested
// with its smallest installed model.
func cliSelftest(args []string) int {
	want := map[string]string{}
	var order []string
	for _, a := range args {
		parts := strings.SplitN(a, "=", 2)
		want[parts[0]] = ""
		if len(parts) == 2 {
			want[parts[0]] = parts[1]
		}
		order = append(order, parts[0])
	}
	if len(order) == 0 {
		for _, rt := range allRuntimes {
			order = append(order, rt.ID())
		}
	}
	failed := 0
	for _, id := range order {
		rt := runtimeByID(id)
		if rt == nil {
			fmt.Printf("FAIL  %-9s unknown runtime\n", id)
			failed++
			continue
		}
		if err := rt.Ready(); err != nil {
			fmt.Printf("FAIL  %-9s not ready: %s\n", id, oneLine(err.Error(), 160))
			failed++
			continue
		}
		inst, err := rt.Installed()
		if err != nil {
			fmt.Printf("FAIL  %-9s cannot list models: %v\n", id, err)
			failed++
			continue
		}
		model := want[id]
		if model == "" {
			model = smallestInstalled(inst)
		} else {
			model = resolveModel(model, inst)
		}
		if _, ok := inst[model]; !ok || model == "" {
			fmt.Printf("FAIL  %-9s no model installed to test (%q)\n", id, model)
			failed++
			continue
		}
		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
		t0 := time.Now()
		if err := rt.EnsureLoaded(ctx, model); err != nil {
			cancel()
			fmt.Printf("FAIL  %-9s %s: loading failed: %s\n", id, model, oneLine(err.Error(), 200))
			failed++
			continue
		}
		loaded := time.Since(t0)
		base, err := rt.BaseURL(model)
		var reply string
		if err == nil {
			t1 := time.Now()
			reply, err = askOnce(ctx, base, model)
			if err == nil {
				fmt.Printf("PASS  %-9s %-32s load %.1fs, answer in %.1fs: %q\n", id, model, loaded.Seconds(), time.Since(t1).Seconds(), oneLine(reply, 60))
			}
		}
		if err != nil {
			fmt.Printf("FAIL  %-9s %s: %s\n", id, model, oneLine(err.Error(), 200))
			failed++
		}
		rt.Unload(model) // leave the memory free again
		cancel()
	}
	if failed > 0 {
		fmt.Printf("%d runtime(s) failed the self-test\n", failed)
		return 1
	}
	fmt.Println("all runtimes passed the self-test")
	return 0
}