package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"
)

// A modelRuntime is something that can hold and serve models (Ollama, LM Studio, vLLM).
// They share the machine's (unified) memory, so the Model Pool looks at all of them together.

type progressFn func(status string, total, done int64)

type modelRuntime interface {
	ID() string
	// Ready returns nil when the runtime is installed and reachable, otherwise why not.
	Ready() error
	Installed() (map[string]int64, error) // model -> bytes on disk
	Loaded() (map[string]int64, error)    // model -> bytes in memory
	// EnsureLoaded makes sure the model is in memory. Runtimes that load on demand do nothing.
	EnsureLoaded(ctx context.Context, name string) error
	Unload(name string) error
	BaseURL(name string) (string, error) // where OpenAI-style requests for this model go
	Pull(ctx context.Context, name string, progress progressFn) error
	Delete(name string) error
}

var allRuntimes = []modelRuntime{ollamaRuntime{}, newLMStudioRuntime(), newVLLMRuntime()}

func runtimeByID(id string) modelRuntime {
	for _, r := range allRuntimes {
		if r.ID() == id {
			return r
		}
	}
	return nil
}

// multiBackend lets the pool treat all ready runtimes as one backend.
type multiBackend struct {
	mu     sync.Mutex
	owners map[string]modelRuntime // model -> runtime, refreshed by installed()
}

func (m *multiBackend) loaded() (map[string]int64, error) {
	out := map[string]int64{}
	var firstErr error
	any := false
	for _, r := range allRuntimes {
		if r.Ready() != nil {
			continue
		}
		l, err := r.Loaded()
		if err != nil {
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		any = true
		for k, v := range l {
			out[k] = v
		}
	}
	if !any && firstErr != nil {
		return nil, firstErr
	}
	return out, nil
}

func (m *multiBackend) installed() (map[string]int64, error) {
	out := map[string]int64{}
	owners := map[string]modelRuntime{}
	var firstErr error
	any := false
	for _, r := range allRuntimes {
		if r.Ready() != nil {
			continue
		}
		inst, err := r.Installed()
		if err != nil {
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		any = true
		for k, v := range inst {
			if _, taken := owners[k]; taken {
				continue // first runtime in the list wins a name clash
			}
			out[k] = v
			owners[k] = r
		}
	}
	if !any && firstErr != nil {
		return nil, firstErr
	}
	m.mu.Lock()
	m.owners = owners
	m.mu.Unlock()
	return out, nil
}

func (m *multiBackend) owner(name string) modelRuntime {
	m.mu.Lock()
	r := m.owners[name]
	m.mu.Unlock()
	if r != nil {
		return r
	}
	m.installed() // refresh and retry once
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.owners[name]
}

func (m *multiBackend) unload(name string) error {
	if r := m.owner(name); r != nil {
		return r.Unload(name)
	}
	// not in the owner table (e.g. just loaded): try every runtime that has it in memory
	for _, r := range allRuntimes {
		if r.Ready() != nil {
			continue
		}
		if l, err := r.Loaded(); err == nil {
			if _, ok := l[name]; ok {
				return r.Unload(name)
			}
		}
	}
	return fmt.Errorf("model %s not found in any runtime", name)
}

func (m *multiBackend) memory() (uint64, uint64) {
	total, avail, _ := budgetMemory()
	return total, avail
}

// ---- Ollama as a runtime (loads on demand, one daemon for all models) ----

type ollamaRuntime struct{}

func (ollamaRuntime) ID() string { return "ollama" }

func (ollamaRuntime) Ready() error {
	var v struct {
		Version string `json:"version"`
	}
	return ollamaBackend{}.get("/api/version", &v)
}

func (ollamaRuntime) Installed() (map[string]int64, error) {
	inst, err := ollamaBackend{}.installed()
	if err != nil {
		return nil, err
	}
	return ollamaDropOffline(loadNodeConfig(), inst), nil // models on a disk that is not connected are not usable
}
func (ollamaRuntime) Loaded() (map[string]int64, error)          { return ollamaBackend{}.loaded() }
func (ollamaRuntime) EnsureLoaded(context.Context, string) error { return nil }
func (ollamaRuntime) Unload(name string) error                   { return ollamaBackend{}.unload(name) }
func (ollamaRuntime) BaseURL(string) (string, error)             { return loadNodeConfig().OllamaURL, nil }

func (ollamaRuntime) Pull(ctx context.Context, name string, progress progressFn) error {
	cfg := loadNodeConfig()
	b, _ := json.Marshal(map[string]any{"name": name, "stream": true})
	req, _ := http.NewRequestWithContext(ctx, "POST", cfg.OllamaURL+"/api/pull", bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("ollama %s", resp.Status)
	}
	return scanNDJSON(resp.Body, func(l map[string]any) error {
		if e, ok := l["error"].(string); ok && e != "" {
			return fmt.Errorf("%s", e)
		}
		st, _ := l["status"].(string)
		tot, _ := l["total"].(float64)
		done, _ := l["completed"].(float64)
		progress(st, int64(tot), int64(done))
		return nil
	})
}

func (ollamaRuntime) Delete(name string) error {
	cfg := loadNodeConfig()
	blobs, _ := ollamaBlobs(cfg.ModelsDir, name)
	b, _ := json.Marshal(map[string]string{"name": name})
	req, _ := http.NewRequest("DELETE", cfg.OllamaURL+"/api/delete", bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("delete failed: %s", resp.Status)
	}
	ollamaForgetAfterDelete(cfg, name, blobs)
	return nil
}