package main

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

// ollamaBackend lets the pool talk to the local Ollama.
type ollamaBackend struct{}

func (ollamaBackend) get(path string, out any) error {
	cfg := loadNodeConfig()
	c := &http.Client{Timeout: 5 * time.Second}
	resp, err := c.Get(cfg.OllamaURL + path)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("ollama %s", resp.Status)
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

func (o ollamaBackend) loaded() (map[string]int64, error) {
	var ps struct {
		Models []struct {
			Name string `json:"name"`
			Size int64  `json:"size"`
		} `json:"models"`
	}
	if err := o.get("/api/ps", &ps); err != nil {
		return nil, err
	}
	out := map[string]int64{}
	for _, m := range ps.Models {
		out[m.Name] = m.Size
	}
	return out, nil
}

func (o ollamaBackend) installed() (map[string]int64, error) {
	var tags struct {
		Models []struct {
			Name string `json:"name"`
			Size int64  `json:"size"`
		} `json:"models"`
	}
	if err := o.get("/api/tags", &tags); err != nil {
		return nil, err
	}
	out := map[string]int64{}
	for _, m := range tags.Models {
		out[m.Name] = m.Size
	}
	return out, nil
}

func (ollamaBackend) unload(name string) error {
	cfg := loadNodeConfig()
	b, _ := json.Marshal(map[string]any{"model": name, "keep_alive": 0})
	resp, err := (&http.Client{Timeout: 30 * time.Second}).Post(cfg.OllamaURL+"/api/generate", "application/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	resp.Body.Close()
	return nil
}

func (ollamaBackend) memory() (uint64, uint64) {
	total, avail, _ := budgetMemory()
	return total, avail
}

var theMulti = &multiBackend{}

var thePool = func() *pool {
	p := newPool(theMulti)
	if gb := loadNodeConfig().ReserveGB; gb > 0 {
		p.reserveBytes = uint64(gb) << 30
	}
	return p
}()

// ---- inference passthrough (server -> agent -> Ollama's OpenAI-compatible API) ----

var allowedInferPaths = map[string]bool{
	"/v1/chat/completions": true,
	"/v1/completions":      true,
	"/v1/embeddings":       true,
}

var inferCancels = struct {
	sync.Mutex
	m map[string]context.CancelFunc
}{m: map[string]context.CancelFunc{}}

func cancelAllInfer() {
	inferCancels.Lock()
	cs := make([]context.CancelFunc, 0, len(inferCancels.m))
	for _, c := range inferCancels.m {
		cs = append(cs, c)
	}
	inferCancels.Unlock()
	for _, c := range cs {
		c()
	}
}

func cancelInfer(id string) {
	inferCancels.Lock()
	c := inferCancels.m[id]
	inferCancels.Unlock()
	if c != nil {
		c()
	}
}

func handleInfer(m Msg, send func(Msg)) {
	fail := func(status int, msg string) {
		b, _ := json.Marshal(map[string]any{"error": map[string]any{"message": msg, "type": "wxch_node_error"}})
		send(Msg{T: "infer_start", ID: m.ID, Status: status, Method: "application/json"})
		send(Msg{T: "infer_chunk", ID: m.ID, D: base64.StdEncoding.EncodeToString(b)})
		send(Msg{T: "infer_end", ID: m.ID})
	}
	if !allowedInferPaths[m.Method] {
		fail(404, "unsupported path")
		return
	}
	if ok, why := guardAccepting(); !ok {
		fail(503, "this machine is not available right now ("+why+")")
		return
	}
	// at the CPU/GPU load limit the request waits for the load to fall; it is refused only when that takes too long
	if why := waitForLoad(context.Background(), inferLoadWait); why != "" {
		fail(503, why)
		return
	}
	var req struct {
		Model string `json:"model"`
	}
	if json.Unmarshal(m.Params, &req) != nil || req.Model == "" {
		fail(400, "request must include a model")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
	defer cancel()
	inferCancels.Lock()
	inferCancels.m[m.ID] = cancel
	inferCancels.Unlock()
	defer func() {
		inferCancels.Lock()
		delete(inferCancels.m, m.ID)
		inferCancels.Unlock()
	}()

	// use the exact installed name from here on (runtimes differ: "qwen2.5:0.5b", "qwen3-0.6b", "org/model")
	inst, _ := theMulti.installed()
	name := resolveModel(req.Model, inst)
	release, perr := thePool.acquire(ctx, name)
	if perr != nil {
		fail(perr.Status, perr.Msg)
		return
	}
	defer release()

	rt := theMulti.owner(name)
	if rt == nil {
		fail(404, fmt.Sprintf("model %q is not installed on this machine", name))
		return
	}
	if err := rt.EnsureLoaded(ctx, name); err != nil {
		if ctx.Err() != nil {
			fail(503, "request cancelled while the model was loading (the machine was handed back to its owner, or the caller gave up)")
			return
		}
		fail(503, err.Error())
		return
	}
	base, err := rt.BaseURL(name)
	if err != nil {
		fail(503, err.Error())
		return
	}
	body := withModel(m.Params, name)
	hreq, _ := http.NewRequestWithContext(ctx, "POST", base+m.Method, bytes.NewReader(body))
	hreq.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{}).Do(hreq)
	if err != nil {
		if ctx.Err() != nil {
			fail(503, "request cancelled (the machine was handed back to its owner, or the caller gave up)")
			return
		}
		fail(502, "model runtime error: "+err.Error())
		return
	}
	defer resp.Body.Close()
	ct := resp.Header.Get("Content-Type")
	if ct == "" {
		ct = "application/json"
	}
	send(Msg{T: "infer_start", ID: m.ID, Status: resp.StatusCode, Method: ct})
	buf := make([]byte, 16*1024)
	for {
		n, err := resp.Body.Read(buf)
		if n > 0 {
			send(Msg{T: "infer_chunk", ID: m.ID, D: base64.StdEncoding.EncodeToString(buf[:n])})
		}
		if err != nil {
			if err != io.EOF && ctx.Err() == nil {
				send(Msg{T: "infer_end", ID: m.ID, Err: err.Error()})
				return
			}
			break
		}
	}
	send(Msg{T: "infer_end", ID: m.ID})
}

// withModel returns the request body with its "model" field set to name.
func withModel(body []byte, name string) []byte {
	var m map[string]json.RawMessage
	if json.Unmarshal(body, &m) != nil {
		return body
	}
	b, _ := json.Marshal(name)
	m["model"] = b
	out, err := json.Marshal(m)
	if err != nil {
		return body
	}
	return out
}