package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

// vLLM runtime: one Docker container per model (vLLM loads a model at start-up and keeps it).
// Model files are plain Hugging Face snapshots under <vllm_dir>/<org>/<repo>, mounted read-only into the container.
// The container is named wxch-vllm-<slug> and labelled, so the agent can find it again after a restart.

const vllmLabel = "wxch.vllm"

var hfRepoRe = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$`)

type vllmRuntime struct {
	mu       sync.Mutex
	readyAt  time.Time
	readyErr error
	sizes    map[string]int64
	sizesAt  time.Time
	loading  map[string]*sync.Mutex
	pulling  map[string]bool
}

func newVLLMRuntime() *vllmRuntime { return &vllmRuntime{loading: map[string]*sync.Mutex{}} }

func (v *vllmRuntime) ID() string { return "vllm" }

func slugOf(name string) string {
	return strings.NewReplacer("/", "--", ".", "-", "_", "-").Replace(strings.ToLower(name))
}
func containerOf(name string) string { return "wxch-vllm-" + slugOf(name) }

// docker runs the docker CLI. A permission problem gets an actionable message.
func docker(timeout time.Duration, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	return dockerCtx(ctx, args...)
}

func dockerCtx(ctx context.Context, args ...string) (string, error) {
	cmd := exec.CommandContext(ctx, "docker", args...)
	var out, errb bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &errb
	err := cmd.Run()
	if err != nil {
		msg := strings.TrimSpace(errb.String())
		if strings.Contains(msg, "permission denied") {
			msg = "vLLM is not installed yet: " + installHint()
		}
		if msg == "" {
			msg = err.Error()
		}
		return out.String(), fmt.Errorf("%s", msg)
	}
	return out.String(), nil
}

func (v *vllmRuntime) Ready() error {
	v.mu.Lock()
	defer v.mu.Unlock()
	if time.Since(v.readyAt) < 15*time.Second {
		return v.readyErr
	}
	v.readyAt = time.Now()
	cfg := loadNodeConfig()
	v.readyErr = pickControl(cfg).Ready(cfg)
	return v.readyErr
}

func dirSize(root string) int64 {
	var n int64
	filepath.WalkDir(root, func(p string, d os.DirEntry, err error) error {
		if err == nil && !d.IsDir() {
			if fi, e := d.Info(); e == nil {
				n += fi.Size()
			}
		}
		return nil
	})
	return n
}

func (v *vllmRuntime) Installed() (map[string]int64, error) {
	v.mu.Lock()
	defer v.mu.Unlock()
	if v.sizes != nil && time.Since(v.sizesAt) < 30*time.Second {
		return v.sizes, nil
	}
	root := loadNodeConfig().VLLMDir
	out := map[string]int64{}
	orgs, _ := os.ReadDir(root)
	for _, o := range orgs {
		if !o.IsDir() {
			continue
		}
		repos, _ := os.ReadDir(filepath.Join(root, o.Name()))
		for _, r := range repos {
			dir := filepath.Join(root, o.Name(), r.Name())
			real, e := filepath.EvalSymlinks(dir) // a model on the external disk is a link
			if e != nil {
				continue // dangling: the disk is not connected
			}
			if _, err := os.Stat(filepath.Join(real, "config.json")); err == nil && !hasPartial(real) {
				out[o.Name()+"/"+r.Name()] = dirSize(real)
			}
		}
	}
	v.sizes, v.sizesAt = out, time.Now()
	return out, nil
}

func hasPartial(dir string) bool {
	found := false
	filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error {
		if err == nil && !d.IsDir() && strings.HasSuffix(p, ".part") {
			found = true
		}
		return nil
	})
	return found
}

type vllmContainer struct {
	Name  string
	Model string
	Port  int
	Mem   int64
	State string
}

func (v *vllmRuntime) containers() ([]vllmContainer, error) {
	return pickControl(loadNodeConfig()).Containers()
}

func (v *vllmRuntime) Loaded() (map[string]int64, error) {
	cs, err := v.containers()
	if err != nil {
		return nil, err
	}
	out := map[string]int64{}
	for _, c := range cs {
		if c.State == "running" {
			out[c.Model] = c.Mem
		}
	}
	return out, nil
}

func (v *vllmRuntime) BaseURL(name string) (string, error) {
	cs, err := v.containers()
	if err != nil {
		return "", err
	}
	for _, c := range cs {
		if c.Model == name && c.State == "running" {
			return fmt.Sprintf("http://127.0.0.1:%d", c.Port), nil
		}
	}
	return "", fmt.Errorf("vLLM model %s is not running", name)
}

func (v *vllmRuntime) lockFor(name string) *sync.Mutex {
	v.mu.Lock()
	defer v.mu.Unlock()
	m := v.loading[name]
	if m == nil {
		m = &sync.Mutex{}
		v.loading[name] = m
	}
	return m
}

func healthy(base string) bool {
	resp, err := (&http.Client{Timeout: 3 * time.Second}).Get(base + "/health")
	if err != nil {
		return false
	}
	resp.Body.Close()
	return resp.StatusCode == 200
}

// plannedMemory is what we ask vLLM to reserve: weights plus room for the KV cache.
func plannedMemory(weights int64, total uint64, cfg nodeConfig) (bytes int64, util float64) {
	if cfg.VLLMGPUUtil > 0 {
		return int64(cfg.VLLMGPUUtil * float64(total)), cfg.VLLMGPUUtil
	}
	b := int64(float64(weights)*1.15) + int64(cfg.VLLMKVGB)<<30
	if total == 0 {
		return b, 0.5
	}
	u := float64(b) / float64(total)
	if u > 0.85 {
		u = 0.85
	}
	return b, u
}

func (v *vllmRuntime) EnsureLoaded(ctx context.Context, name string) error {
	if !hfRepoRe.MatchString(name) {
		return fmt.Errorf("invalid vLLM model name")
	}
	l := v.lockFor(name)
	l.Lock()
	defer l.Unlock()
	if base, err := v.BaseURL(name); err == nil && healthy(base) {
		return nil
	}
	cfg := loadNodeConfig()
	inst, _ := v.Installed()
	weights, ok := inst[name]
	if !ok {
		return fmt.Errorf("model %s is not downloaded for vLLM", name)
	}
	ctl := pickControl(cfg)
	cs, _ := v.containers()
	used := map[int]bool{}
	for _, c := range cs {
		if c.Model != name {
			used[c.Port] = true
		}
	}
	port := cfg.VLLMBasePort
	for used[port] {
		port++
	}
	// vLLM's utilisation fraction is per GPU: use the size of one GPU on discrete-GPU machines, system RAM otherwise
	total, _, perGPU := budgetMemory()
	if perGPU > 0 {
		total = perGPU
	}
	mem, util := plannedMemory(weights, total, cfg)
	if err := ctl.Start(cfg, name, port, mem, util); err != nil {
		return fmt.Errorf("cannot start vLLM: %v", err)
	}
	base := fmt.Sprintf("http://127.0.0.1:%d", port)
	deadline := time.Now().Add(20 * time.Minute)
	for {
		if healthy(base) {
			return nil
		}
		select {
		case <-ctx.Done():
			ctl.Stop(name)
			return fmt.Errorf("loading was cancelled")
		case <-time.After(3 * time.Second):
		}
		if !ctl.Running(name) {
			logs := ctl.Logs(name, 25)
			ctl.Stop(name)
			return fmt.Errorf("vLLM exited while loading %s: %s", name, tailText(logs, 900))
		}
		if time.Now().After(deadline) {
			ctl.Stop(name)
			return fmt.Errorf("vLLM did not become ready within 20 minutes")
		}
	}
}

func tailText(s string, n int) string {
	s = strings.TrimSpace(s)
	if len(s) > n {
		s = s[len(s)-n:]
	}
	return s
}

func (v *vllmRuntime) Unload(name string) error {
	return pickControl(loadNodeConfig()).Stop(name)
}

func (v *vllmRuntime) Delete(name string) error {
	if !hfRepoRe.MatchString(name) {
		return fmt.Errorf("invalid model name")
	}
	if l, err := v.Loaded(); err == nil {
		if _, busy := l[name]; busy {
			return fmt.Errorf("model is loaded in memory; unload it first")
		}
	}
	root, _ := filepath.Abs(loadNodeConfig().VLLMDir)
	dir, _ := filepath.Abs(filepath.Join(root, name))
	if !strings.HasPrefix(dir, root+string(os.PathSeparator)) {
		return fmt.Errorf("invalid path")
	}
	v.mu.Lock()
	v.sizes = nil
	v.mu.Unlock()
	return vllmRemoveUnit(loadNodeConfig(), dir)
}

// ---- Hugging Face download ----

// hfBase is a variable so tests can point it at a fake server.
var hfBase = "https://huggingface.co"

func hfToken() string {
	if t := strings.TrimSpace(os.Getenv("HF_TOKEN")); t != "" {
		return t
	}
	if b, err := os.ReadFile("/etc/wxch/hf_token"); err == nil {
		return strings.TrimSpace(string(b))
	}
	return ""
}

func hfRequest(ctx context.Context, url string, rangeFrom int64) (*http.Response, error) {
	req, _ := http.NewRequestWithContext(ctx, "GET", url, nil)
	if t := hfToken(); t != "" {
		req.Header.Set("Authorization", "Bearer "+t)
	}
	if rangeFrom > 0 {
		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", rangeFrom))
	}
	return (&http.Client{}).Do(req)
}

func (v *vllmRuntime) Pull(ctx context.Context, name string, progress progressFn) error {
	return v.pullWith(ctx, loadNodeConfig(), name, progress)
}

func (v *vllmRuntime) pullWith(ctx context.Context, cfg nodeConfig, name string, progress progressFn) error {
	if !hfRepoRe.MatchString(name) {
		return fmt.Errorf("vLLM models are Hugging Face ids like org/name")
	}
	// All callers share this runtime, including separate browser sessions.
	v.mu.Lock()
	if v.pulling[name] {
		v.mu.Unlock()
		return fmt.Errorf("model %s is already downloading; wait for the current download", name)
	}
	if v.pulling == nil {
		v.pulling = make(map[string]bool)
	}
	v.pulling[name] = true
	v.mu.Unlock()
	defer func() {
		v.mu.Lock()
		delete(v.pulling, name)
		v.mu.Unlock()
	}()
	progress("looking up "+name, 0, 0)
	resp, err := hfRequest(ctx, hfBase+"/api/models/"+name+"?blobs=true", 0)
	if err != nil {
		return err
	}
	var info struct {
		SHA      string `json:"sha"`
		Gated    any    `json:"gated"`
		Siblings []struct {
			Name string `json:"rfilename"`
			Size int64  `json:"size"`
		} `json:"siblings"`
	}
	err = json.NewDecoder(resp.Body).Decode(&info)
	status := resp.StatusCode
	resp.Body.Close()
	if status == 401 || status == 403 {
		return fmt.Errorf("Hugging Face refused access to %s (gated model? put a token in /etc/wxch/hf_token)", name)
	}
	if status != 200 || err != nil || len(info.Siblings) == 0 {
		return fmt.Errorf("model %s not found on Hugging Face (HTTP %d)", name, status)
	}
	type file struct {
		name string
		size int64
	}
	var files []file
	var total int64
	for _, s := range info.Siblings {
		l := strings.ToLower(s.Name)
		if strings.HasPrefix(s.Name, ".") || strings.HasSuffix(l, ".md") || strings.Contains(s.Name, "..") || strings.HasPrefix(s.Name, "/") {
			continue
		}
		files = append(files, file{s.Name, s.Size})
		total += s.Size
	}
	root := filepath.Join(cfg.VLLMDir, name)
	if free, _, err := diskFree(cfg.VLLMDir); err == nil && uint64(total)+uint64(cfg.MinFreeGB)<<30 > free {
		return fmt.Errorf("not enough free disk space: need %.1f GB (+%d GB spare), have %.1f GB", float64(total)/1e9, cfg.MinFreeGB, float64(free)/1e9)
	}
	var done int64
	for _, f := range files {
		dest := filepath.Join(root, filepath.FromSlash(f.name))
		if fi, err := os.Stat(dest); err == nil && fi.Size() == f.size {
			done += f.size
			progress("have "+f.name, total, done)
			continue
		}
		if err := os.MkdirAll(filepath.Dir(dest), 0755); err != nil {
			return err
		}
		part := dest + ".part"
		var have int64
		if fi, err := os.Stat(part); err == nil {
			have = fi.Size()
		}
		// An oversized partial cannot be resumed: Range would return 416 forever.
		// Request the full file; the 200 response below replaces only this partial.
		if have > f.size {
			have = 0
		}
		r, err := hfRequest(ctx, hfBase+"/"+name+"/resolve/"+info.SHA+"/"+f.name, have)
		if err != nil {
			return err
		}
		if r.StatusCode == 416 { // the .part is already complete
			r.Body.Close()
			have = f.size
		} else if r.StatusCode != 200 && r.StatusCode != 206 {
			r.Body.Close()
			return fmt.Errorf("download of %s failed: HTTP %d", f.name, r.StatusCode)
		} else {
			flag := os.O_CREATE | os.O_WRONLY | os.O_APPEND
			if r.StatusCode == 200 { // server ignored Range: start over
				flag = os.O_CREATE | os.O_WRONLY | os.O_TRUNC
				have = 0
			}
			out, err := os.OpenFile(part, flag, 0644)
			if err != nil {
				r.Body.Close()
				return err
			}
			cur := done + have
			buf := make([]byte, 1<<20)
			for {
				n, rerr := r.Body.Read(buf)
				if n > 0 {
					if _, werr := out.Write(buf[:n]); werr != nil {
						out.Close()
						r.Body.Close()
						return werr
					}
					cur += int64(n)
					progress("downloading "+f.name, total, cur)
				}
				if rerr == io.EOF {
					break
				}
				if rerr != nil {
					out.Close()
					r.Body.Close()
					return fmt.Errorf("download of %s interrupted: %v (run it again to resume)", f.name, rerr)
				}
			}
			out.Close()
			r.Body.Close()
		}
		if fi, err := os.Stat(part); err != nil || fi.Size() != f.size {
			return fmt.Errorf("%s has the wrong size after download (expected %d bytes; run it again)", f.name, f.size)
		}
		if err := os.Rename(part, dest); err != nil {
			return err
		}
		done += f.size
	}
	v.mu.Lock()
	v.sizes = nil
	v.mu.Unlock()
	progress("complete", total, total)
	return nil
}

// scanNDJSON calls fn for each JSON line of r.
func scanNDJSON(r io.Reader, fn func(map[string]any) error) error {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 1<<20), 1<<20)
	for sc.Scan() {
		var m map[string]any
		if json.Unmarshal(sc.Bytes(), &m) != nil {
			continue
		}
		if err := fn(m); err != nil {
			return err
		}
	}
	return sc.Err()
}