package main

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

// LM Studio (llmster, headless) runtime, driven through its REST API:
//   GET  /api/v1/models                      list (with loaded instances)
//   POST /api/v1/models/load                 load a model with our settings
//   POST /api/v1/models/unload               unload an instance
//   POST /api/v1/models/download             start a download -> job id
//   GET  /api/v1/models/download/status/:id  progress
// Inference goes to the OpenAI-compatible /v1/* endpoints of the same server.

type lmstudioRuntime struct {
	mu       sync.Mutex
	readyAt  time.Time
	readyErr error
	loadMu   map[string]*sync.Mutex
}

func newLMStudioRuntime() *lmstudioRuntime { return &lmstudioRuntime{loadMu: map[string]*sync.Mutex{}} }

func (*lmstudioRuntime) ID() string { return "lmstudio" }

type lmModel struct {
	Type            string `json:"type"`
	Key             string `json:"key"`
	DisplayName     string `json:"display_name"`
	SizeBytes       int64  `json:"size_bytes"`
	LoadedInstances []struct {
		ID string `json:"id"`
	} `json:"loaded_instances"`
}

func lmHTTP(ctx context.Context, cfg nodeConfig, method, path string, body any, timeout time.Duration) (*http.Response, error) {
	var rd io.Reader
	if body != nil {
		b, _ := json.Marshal(body)
		rd = bytes.NewReader(b)
	}
	base := cfg.LMStudioURL
	if lmTestURL != "" {
		base = lmTestURL
	}
	req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(base, "/")+path, rd)
	if err != nil {
		return nil, err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if tok := lmToken(cfg); tok != "" {
		req.Header.Set("Authorization", "Bearer "+tok)
	}
	return (&http.Client{Timeout: timeout}).Do(req)
}

// lmTestURL overrides the server URL in tests.
var lmTestURL string

func lmToken(cfg nodeConfig) string {
	if cfg.LMStudioToken != "" {
		return cfg.LMStudioToken
	}
	return strings.TrimSpace(os.Getenv("LM_API_TOKEN"))
}

func lmList(cfg nodeConfig) ([]lmModel, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	resp, err := lmHTTP(ctx, cfg, "GET", "/api/v1/models", nil, 8*time.Second)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 300))
		return nil, fmt.Errorf("LM Studio %s %s", resp.Status, strings.TrimSpace(string(b)))
	}
	var out struct {
		Models []lmModel `json:"models"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, err
	}
	var llms []lmModel
	for _, m := range out.Models {
		if m.Type == "" || m.Type == "llm" {
			llms = append(llms, m)
		}
	}
	return llms, nil
}

func (l *lmstudioRuntime) Ready() error {
	l.mu.Lock()
	defer l.mu.Unlock()
	if time.Since(l.readyAt) < 10*time.Second {
		return l.readyErr
	}
	l.readyAt = time.Now()
	cfg := loadNodeConfig()
	if _, err := lmList(cfg); err != nil {
		if bin := lmBin(cfg); bin == "" {
			l.readyErr = fmt.Errorf("LM Studio is not installed: %s", installHint())
		} else {
			l.readyErr = fmt.Errorf("LM Studio server is not running (%v)", shortErr(err))
		}
		return l.readyErr
	}
	l.readyErr = nil
	return nil
}

func shortErr(err error) string {
	s := err.Error()
	if len(s) > 120 {
		s = s[:120]
	}
	return s
}

// lmBin finds the lms command line tool, or "" if LM Studio is not installed for this user.
func lmBin(cfg nodeConfig) string {
	cands := []string{cfg.LMStudioBin}
	if home, err := os.UserHomeDir(); err == nil {
		cands = append(cands, home+"/.lmstudio/bin/lms")
	}
	if p, err := exec.LookPath("lms"); err == nil {
		cands = append(cands, p)
	}
	for _, c := range cands {
		if c == "" {
			continue
		}
		if fi, err := os.Stat(c); err == nil && !fi.IsDir() {
			return c
		}
	}
	return ""
}

func (*lmstudioRuntime) Installed() (map[string]int64, error) {
	ms, err := lmList(loadNodeConfig())
	if err != nil {
		return nil, err
	}
	out := map[string]int64{}
	for _, m := range ms {
		out[m.Key] = m.SizeBytes
	}
	return out, nil
}

func (*lmstudioRuntime) Loaded() (map[string]int64, error) {
	ms, err := lmList(loadNodeConfig())
	if err != nil {
		return nil, err
	}
	out := map[string]int64{}
	for _, m := range ms {
		if len(m.LoadedInstances) > 0 {
			out[m.Key] = m.SizeBytes + m.SizeBytes/6 // weights plus a share for the KV cache and buffers
		}
	}
	return out, nil
}

func (*lmstudioRuntime) BaseURL(string) (string, error) {
	if lmTestURL != "" {
		return lmTestURL, nil
	}
	return strings.TrimRight(loadNodeConfig().LMStudioURL, "/"), nil
}

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

func (l *lmstudioRuntime) EnsureLoaded(ctx context.Context, name string) error {
	lk := l.lockFor(name)
	lk.Lock()
	defer lk.Unlock()
	cfg := loadNodeConfig()
	if loaded, err := l.Loaded(); err == nil {
		if _, ok := loaded[name]; ok {
			return nil
		}
	}
	resp, err := lmHTTP(ctx, cfg, "POST", "/api/v1/models/load", map[string]any{"model": name, "context_length": cfg.LMStudioCtx}, 20*time.Minute)
	if err != nil {
		return fmt.Errorf("cannot load %s in LM Studio: %v", name, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 400))
		return fmt.Errorf("LM Studio could not load %s: %s %s", name, resp.Status, strings.TrimSpace(string(b)))
	}
	return nil
}

func (*lmstudioRuntime) Unload(name string) error {
	cfg := loadNodeConfig()
	ms, err := lmList(cfg)
	if err != nil {
		return err
	}
	for _, m := range ms {
		if m.Key != name {
			continue
		}
		for _, inst := range m.LoadedInstances {
			ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
			resp, err := lmHTTP(ctx, cfg, "POST", "/api/v1/models/unload", map[string]string{"instance_id": inst.ID}, 60*time.Second)
			cancel()
			if err != nil {
				return err
			}
			resp.Body.Close()
		}
		return nil
	}
	return fmt.Errorf("model %s not found in LM Studio", name)
}

// Pull accepts a catalog id ("ibm/granite-4-micro"), a Hugging Face repo ("org/name") or
// "org/name@Q4_K_M" to pick a quantization of a Hugging Face repo.
func (*lmstudioRuntime) Pull(ctx context.Context, name string, progress progressFn) error {
	cfg := loadNodeConfig()
	req := map[string]any{"model": name}
	if i := strings.Index(name, "@"); i > 0 {
		req["model"] = "https://huggingface.co/" + name[:i]
		req["quantization"] = name[i+1:]
	}
	resp, err := lmHTTP(ctx, cfg, "POST", "/api/v1/models/download", req, 60*time.Second)
	if err != nil {
		return err
	}
	var job struct {
		JobID  string `json:"job_id"`
		Status string `json:"status"`
		Total  int64  `json:"total_size_bytes"`
	}
	err = json.NewDecoder(resp.Body).Decode(&job)
	code := resp.StatusCode
	resp.Body.Close()
	if code != 200 || err != nil {
		return fmt.Errorf("LM Studio refused the download (HTTP %d)", code)
	}
	if job.Status == "already_downloaded" || job.Status == "completed" {
		progress("complete", job.Total, job.Total)
		return nil
	}
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
		r, err := lmHTTP(ctx, cfg, "GET", "/api/v1/models/download/status/"+job.JobID, nil, 30*time.Second)
		if err != nil {
			return err
		}
		var st struct {
			Status     string `json:"status"`
			Total      int64  `json:"total_size_bytes"`
			Downloaded int64  `json:"downloaded_bytes"`
		}
		json.NewDecoder(r.Body).Decode(&st)
		r.Body.Close()
		switch st.Status {
		case "completed":
			progress("complete", st.Total, st.Total)
			return nil
		case "failed":
			return fmt.Errorf("LM Studio reported the download failed")
		default:
			progress("downloading", st.Total, st.Downloaded)
		}
	}
}

// lmFile is one entry of `lms ls --json`: the model key and its file path relative to the models folder.
type lmFile struct {
	ModelKey string `json:"modelKey"`
	Path     string `json:"path"`
}

// lmListFiles asks the lms CLI where each model lives. A variable so tests can replace it.
var lmListFiles = func(cfg nodeConfig) ([]lmFile, error) {
	bin := lmBin(cfg)
	if bin == "" {
		return nil, fmt.Errorf("lms command not found")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	out, err := exec.CommandContext(ctx, bin, "ls", "--json").Output()
	if err != nil {
		return nil, fmt.Errorf("lms ls failed: %v", err)
	}
	var files []lmFile
	if err := json.Unmarshal(out, &files); err != nil {
		return nil, err
	}
	return files, nil
}

// lmDeleteWait is how many seconds Delete waits for LM Studio to forget the model (tests set it to 0).
var lmDeleteWait = 20

// lmRootOverride lets tests point at a temp folder.
var lmRootOverride string

func lmModelsRoot(cfg nodeConfig) string {
	if lmRootOverride != "" {
		return lmRootOverride
	}
	if cfg.LMStudioModels != "" {
		return cfg.LMStudioModels
	}
	home, _ := os.UserHomeDir()
	return filepath.Join(home, ".lmstudio", "models")
}

// Delete removes the model's file(s) from the models folder (LM Studio has no delete command or endpoint).
func (l *lmstudioRuntime) Delete(name string) error {
	cfg := loadNodeConfig()
	if loaded, err := l.Loaded(); err == nil {
		if _, busy := loaded[name]; busy {
			return fmt.Errorf("model is loaded in memory; unload it first")
		}
	}
	files, err := lmListFiles(cfg)
	if err != nil {
		return err
	}
	rel := ""
	for _, f := range files {
		if f.ModelKey == name {
			rel = f.Path
			break
		}
	}
	if rel == "" {
		return fmt.Errorf("model %s not found in LM Studio", name)
	}
	root, err := filepath.EvalSymlinks(lmModelsRoot(cfg))
	if err != nil {
		return err
	}
	file := filepath.Join(root, filepath.FromSlash(rel))
	if !strings.HasPrefix(file, root+string(os.PathSeparator)) || strings.Contains(rel, "..") {
		return fmt.Errorf("refusing to delete outside the models folder")
	}
	if err := os.Remove(file); err != nil {
		return err
	}
	// remove the model's folder once no weights are left (a leftover mmproj/vision file alone is useless)
	dir := filepath.Dir(file)
	if dir != root {
		left, _ := os.ReadDir(dir)
		onlyAux := true
		for _, e := range left {
			if strings.HasSuffix(strings.ToLower(e.Name()), ".gguf") && !strings.HasPrefix(strings.ToLower(e.Name()), "mmproj") {
				onlyAux = false
			}
		}
		if onlyAux {
			os.RemoveAll(dir)
		}
	}
	// LM Studio notices the removed file after a few seconds; wait so the caller does not see the stale entry
	for i := 0; i < lmDeleteWait; i++ {
		inst, err := l.Installed()
		if err != nil {
			break
		}
		if _, still := inst[name]; !still {
			break
		}
		time.Sleep(time.Second)
	}
	return nil
}