package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

// vllmControl is how the vLLM runtime starts and stops containers. There are two implementations:
//   - runnerControl talks to the root-owned wxch-runner service over a unix socket (the default and the
//     safe option: the agent itself has no Docker access);
//   - directControl runs the docker CLI as the agent user (only works if that user is in the docker group,
//     which is equivalent to root, so it is a fallback for machines where that was chosen on purpose).
type vllmControl interface {
	Ready(cfg nodeConfig) error
	Containers() ([]vllmContainer, error)
	Start(cfg nodeConfig, model string, port int, mem int64, util float64) error
	Running(model string) bool
	Logs(model string, tail int) string
	Stop(model string) error
}

const defaultRunnerSocket = "/run/wxch/runner.sock"

func runnerSocket(cfg nodeConfig) string {
	if cfg.VLLMRunnerSocket != "" {
		return cfg.VLLMRunnerSocket
	}
	return defaultRunnerSocket
}

// controlOverride lets tests inject a fake.
var controlOverride vllmControl

func pickControl(cfg nodeConfig) vllmControl {
	if controlOverride != nil {
		return controlOverride
	}
	if _, err := os.Stat(runnerSocket(cfg)); err == nil {
		return runnerControl{sock: runnerSocket(cfg)}
	}
	return directControl{}
}

// ---- runner (unix socket) ----

type runnerControl struct{ sock string }

type runnerRequest struct {
	Op     string   `json:"op"`
	Model  string   `json:"model,omitempty"`
	Port   int      `json:"port,omitempty"`
	Mem    int64    `json:"mem,omitempty"`
	Util   float64  `json:"util,omitempty"`
	MaxLen int      `json:"max_len,omitempty"`
	Args   []string `json:"args,omitempty"`
	Tail   int      `json:"tail,omitempty"`
	Engine string   `json:"engine,omitempty"`
}

func (r runnerControl) call(req runnerRequest, timeout time.Duration) (json.RawMessage, error) {
	c, err := net.DialTimeout("unix", r.sock, 5*time.Second)
	if err != nil {
		return nil, fmt.Errorf("vLLM runner is not reachable (%v): %s", err, installHint())
	}
	defer c.Close()
	c.SetDeadline(time.Now().Add(timeout))
	b, _ := json.Marshal(req)
	if _, err := c.Write(append(b, '\n')); err != nil {
		return nil, err
	}
	line, err := bufio.NewReaderSize(c, 1<<20).ReadBytes('\n')
	if err != nil && len(line) == 0 {
		return nil, fmt.Errorf("no answer from the vLLM runner: %v", err)
	}
	var resp struct {
		OK   bool            `json:"ok"`
		Err  string          `json:"err"`
		Data json.RawMessage `json:"data"`
	}
	if err := json.Unmarshal(line, &resp); err != nil {
		return nil, err
	}
	if !resp.OK {
		return nil, fmt.Errorf("%s", resp.Err)
	}
	return resp.Data, nil
}

func (r runnerControl) Ready(nodeConfig) error {
	_, err := r.call(runnerRequest{Op: "status"}, 20*time.Second)
	return err
}

func (r runnerControl) Containers() ([]vllmContainer, error) {
	data, err := r.call(runnerRequest{Op: "ps"}, 20*time.Second)
	if err != nil {
		return nil, err
	}
	var cs []struct {
		Name, Model, State string
		Port               int
		Mem                int64
	}
	if err := json.Unmarshal(data, &cs); err != nil {
		return nil, err
	}
	out := make([]vllmContainer, 0, len(cs))
	for _, c := range cs {
		out = append(out, vllmContainer{Name: c.Name, Model: c.Model, Port: c.Port, Mem: c.Mem, State: c.State})
	}
	return out, nil
}

func (r runnerControl) Start(cfg nodeConfig, model string, port int, mem int64, util float64) error {
	_, err := r.call(runnerRequest{Op: "start", Model: model, Port: port, Mem: mem, Util: util, MaxLen: cfg.VLLMMaxLen, Args: cfg.VLLMExtraArgs}, 6*time.Minute)
	return err
}

func (r runnerControl) Running(model string) bool {
	data, err := r.call(runnerRequest{Op: "running", Model: model}, 20*time.Second)
	if err != nil {
		return false
	}
	var running bool
	json.Unmarshal(data, &running)
	return running
}

func (r runnerControl) Logs(model string, tail int) string {
	data, err := r.call(runnerRequest{Op: "logs", Model: model, Tail: tail}, 30*time.Second)
	if err != nil {
		return err.Error()
	}
	var s string
	json.Unmarshal(data, &s)
	return s
}

func (r runnerControl) Stop(model string) error {
	_, err := r.call(runnerRequest{Op: "stop", Model: model}, 90*time.Second)
	return err
}

// ---- direct docker CLI (fallback) ----

type directControl struct{}

func (directControl) Ready(cfg nodeConfig) error {
	if _, err := docker(6*time.Second, "info", "--format", "{{.ServerVersion}}"); err != nil {
		return err
	}
	if _, err := docker(6*time.Second, "image", "inspect", "--format", "{{.Id}}", cfg.VLLMImage); err != nil {
		return fmt.Errorf("vLLM image %s is not pulled yet: %s", cfg.VLLMImage, installHint())
	}
	return nil
}

func (directControl) Containers() ([]vllmContainer, error) {
	out, err := docker(8*time.Second, "ps", "-a", "--filter", "label="+vllmLabel,
		"--format", `{{.Names}}|{{.Label "wxch.model"}}|{{.Label "wxch.port"}}|{{.Label "wxch.mem"}}|{{.State}}`)
	if err != nil {
		return nil, err
	}
	var cs []vllmContainer
	for _, line := range strings.Split(strings.TrimSpace(out), "\n") {
		f := strings.Split(line, "|")
		if len(f) < 5 {
			continue
		}
		port, _ := strconv.Atoi(f[2])
		mem, _ := strconv.ParseInt(f[3], 10, 64)
		cs = append(cs, vllmContainer{f[0], f[1], port, mem, f[4]})
	}
	return cs, nil
}

func (directControl) Start(cfg nodeConfig, model string, port int, mem int64, util float64) error {
	cname := containerOf(model)
	docker(30*time.Second, "rm", "-f", cname)
	args := []string{"run", "-d", "--name", cname,
		"--label", vllmLabel + "=1", "--label", "wxch.model=" + model, "--label", "wxch.port=" + strconv.Itoa(port), "--label", "wxch.mem=" + strconv.FormatInt(mem, 10),
		"--gpus", "all", "--ipc=host", "--ulimit", "memlock=-1", "--ulimit", "stack=67108864",
		"-p", fmt.Sprintf("127.0.0.1:%d:8000", port),
		"-v", cfg.VLLMDir + ":/models:ro"}
	if loc, _ := vllmLocation(cfg, model); loc == "external" { // a link to an external disk must resolve inside the container too
		extRoot := cfg.extRoot()
		args = append(args, "-v", extRoot+":"+extRoot+":ro")
	}
	args = append(args, cfg.VLLMImage,
		"vllm", "serve", "/models/"+model, "--served-model-name", model, "--host", "0.0.0.0", "--port", "8000",
		"--max-model-len", strconv.Itoa(cfg.VLLMMaxLen), "--gpu-memory-utilization", fmt.Sprintf("%.2f", util))
	args = append(args, cfg.VLLMExtraArgs...)
	_, err := docker(5*time.Minute, args...)
	return err
}

func (directControl) Running(model string) bool {
	out, _ := docker(5*time.Second, "inspect", "-f", "{{.State.Running}}", containerOf(model))
	return strings.TrimSpace(out) == "true"
}

func (directControl) Logs(model string, tail int) string {
	// docker logs prints the container's stderr on its own stderr, where vLLM writes
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	var out bytes.Buffer
	cmd := exec.CommandContext(ctx, "docker", "logs", "--tail", strconv.Itoa(tail), containerOf(model))
	cmd.Stdout, cmd.Stderr = &out, &out
	cmd.Run()
	return out.String()
}

func (directControl) Stop(model string) error {
	_, err := docker(40*time.Second, "rm", "-f", containerOf(model))
	return err
}