// Package comfy has images made by ComfyUI: a workflow filled in with what
// was asked for, queued, waited on, and its pictures fetched.
//
// ComfyUI speaks no OpenAI. A workflow is a graph of nodes that is queued with
// POST /prompt, is done when GET /history/{id} lists it, and leaves its
// pictures to be read one by one from GET /view. What the graph is depends on
// the model — FLUX schnell is not FLUX dev, and neither is SDXL — so the graph
// is a template the catalog names, with placeholders where a request's
// prompt, size and seed go.
package comfy

import (
	"bytes"
	"context"
	"crypto/rand"
	"embed"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"net/url"
	"slices"
	"strconv"
	"strings"
	"time"
)

//go:embed workflows/*.json
var builtin embed.FS

// DefaultWorkflow is the template used when the catalog names none.
const DefaultWorkflow = "flux-schnell"

// Builtins are the workflows that come with llmgw, by name.
func Builtins() []string {
	entries, _ := builtin.ReadDir("workflows")
	out := make([]string, 0, len(entries))
	for _, e := range entries {
		out = append(out, strings.TrimSuffix(e.Name(), ".json"))
	}
	return out
}

// Template is a workflow in ComfyUI's API format, with placeholders.
//
// A placeholder is a string value "{{name}}", or "{{name|default}}" for one a
// request may leave out. One that is the whole value becomes the value
// itself, a number where the node wants a number; one inside a longer string
// is written into it. The names are those of Params.
type Template struct {
	graph map[string]any
}

// Load reads a template: a built-in's name, or the workflow's JSON itself.
func Load(nameOrJSON string) (*Template, error) {
	s := strings.TrimSpace(nameOrJSON)
	if s == "" {
		s = DefaultWorkflow
	}
	raw := []byte(s)
	if !strings.HasPrefix(s, "{") {
		b, err := builtin.ReadFile("workflows/" + s + ".json")
		if err != nil {
			return nil, fmt.Errorf("no built-in workflow %q (there are %s)", s, strings.Join(Builtins(), ", "))
		}
		raw = b
	}
	var g map[string]any
	if err := json.Unmarshal(raw, &g); err != nil {
		return nil, fmt.Errorf("the workflow is not JSON: %w", err)
	}
	if len(g) == 0 {
		return nil, errors.New("the workflow has no nodes")
	}
	for id, n := range g {
		node, ok := n.(map[string]any)
		if !ok || node["class_type"] == nil {
			// The UI's own save format has a "nodes" list instead, and is
			// not something /prompt accepts.
			return nil, fmt.Errorf("node %q has no class_type: the workflow must be saved in API format", id)
		}
	}
	return &Template{graph: g}, nil
}

// Params are what a request decides. A zero number is one not given, which
// leaves it to the template's default.
type Params struct {
	Checkpoint     string
	Prompt         string
	NegativePrompt string
	Width, Height  int
	N              int
	Seed           int64
	Steps          int
	Guidance       float64
}

func (p Params) value(name string) (any, bool) {
	switch name {
	case "checkpoint":
		return p.Checkpoint, p.Checkpoint != ""
	case "prompt":
		return p.Prompt, true
	case "negative_prompt":
		return p.NegativePrompt, true
	case "width":
		return p.Width, p.Width > 0
	case "height":
		return p.Height, p.Height > 0
	case "n":
		return max(p.N, 1), true
	case "seed":
		return p.Seed, true
	case "steps":
		return p.Steps, p.Steps > 0
	case "guidance":
		return p.Guidance, p.Guidance > 0
	}
	return nil, false
}

// Fill is the workflow for one request.
func (t *Template) Fill(p Params) (map[string]any, error) {
	var missing []string
	var fill func(v any) any
	fill = func(v any) any {
		switch v := v.(type) {
		case map[string]any:
			out := make(map[string]any, len(v))
			for k, x := range v {
				out[k] = fill(x)
			}
			return out
		case []any:
			out := make([]any, len(v))
			for i, x := range v {
				out[i] = fill(x)
			}
			return out
		case string:
			return fillString(v, p, &missing)
		}
		return v
	}
	g := fill(t.graph).(map[string]any)
	if len(missing) > 0 {
		slices.Sort(missing)
		return nil, fmt.Errorf("the workflow wants %s, which the request did not give and the workflow has no default for",
			strings.Join(slices.Compact(missing), ", "))
	}
	return g, nil
}

func fillString(s string, p Params, missing *[]string) any {
	if !strings.Contains(s, "{{") {
		return s
	}
	resolve := func(ph string) (any, bool) {
		name, def, hasDef := strings.Cut(ph, "|")
		name = strings.TrimSpace(name)
		if v, ok := p.value(name); ok {
			return v, true
		}
		if hasDef {
			def = strings.TrimSpace(def)
			if f, err := strconv.ParseFloat(def, 64); err == nil {
				return f, true
			}
			return def, true
		}
		*missing = append(*missing, name)
		return nil, false
	}
	// The whole value: what the node gets is the value, not its text.
	if strings.HasPrefix(s, "{{") && strings.HasSuffix(s, "}}") && strings.Count(s, "{{") == 1 {
		v, _ := resolve(s[2 : len(s)-2])
		return v
	}
	var b strings.Builder
	for {
		i := strings.Index(s, "{{")
		j := strings.Index(s, "}}")
		if i < 0 || j < i {
			b.WriteString(s)
			return b.String()
		}
		b.WriteString(s[:i])
		if v, ok := resolve(s[i+2 : j]); ok {
			fmt.Fprint(&b, v)
		}
		s = s[j+2:]
	}
}

// RandomSeed is a seed for a request that named none. ComfyUI's are 64-bit,
// but JSON numbers are read as doubles on the way: kept to 2^53.
func RandomSeed() int64 {
	n, err := rand.Int(rand.Reader, big.NewInt(1<<53))
	if err != nil {
		return time.Now().UnixNano() & (1<<53 - 1)
	}
	return n.Int64()
}

// Image is one picture made.
type Image struct {
	Data        []byte
	ContentType string
}

// Error is a failure, said as what it means for the request.
type Error struct {
	Msg string
	// Retryable is a failure of this server, which another one might not
	// have: it could not be reached, or it answered 5xx.
	Retryable bool
}

func (e *Error) Error() string { return e.Msg }

// Client queues workflows and waits for them.
type Client struct {
	HTTP *http.Client
	// Poll is how often a queued workflow is asked after.
	Poll time.Duration
}

// Generate runs a workflow on the ComfyUI at base and returns its pictures,
// in the order of the nodes that made them. A context that ends takes the
// workflow out of the queue, or interrupts it if it is running, so a client
// that went away does not leave the card working for nobody.
func (c *Client) Generate(ctx context.Context, base, apiKey string, workflow map[string]any) ([]Image, error) {
	base = strings.TrimRight(base, "/")
	client := make([]byte, 8)
	_, _ = rand.Read(client)
	body, err := json.Marshal(map[string]any{"prompt": workflow, "client_id": "llmgw-" + hex.EncodeToString(client)})
	if err != nil {
		return nil, err
	}
	var queued struct {
		PromptID   string          `json:"prompt_id"`
		Error      json.RawMessage `json:"error"`
		NodeErrors json.RawMessage `json:"node_errors"`
	}
	status, err := c.do(ctx, http.MethodPost, base+"/prompt", apiKey, body, &queued)
	if err != nil {
		return nil, err
	}
	if status != http.StatusOK || queued.PromptID == "" {
		// A workflow ComfyUI will not run: a checkpoint it does not have, a
		// node it does not know. The catalog's to fix, not the client's, and
		// no other replica set up the same way would take it either.
		return nil, &Error{Msg: "ComfyUI refused the workflow: " + refusal(queued.Error, queued.NodeErrors)}
	}
	id := queued.PromptID

	poll := c.Poll
	if poll <= 0 {
		poll = 250 * time.Millisecond
	}
	t := time.NewTicker(poll)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			c.cancel(context.WithoutCancel(ctx), base, apiKey, id)
			return nil, ctx.Err()
		case <-t.C:
		}
		var hist map[string]historyEntry
		if _, err := c.do(ctx, http.MethodGet, base+"/history/"+url.PathEscape(id), apiKey, nil, &hist); err != nil {
			if ctx.Err() != nil {
				continue // said by the select
			}
			return nil, err
		}
		h, ok := hist[id]
		if !ok {
			continue
		}
		if h.Status.StatusStr == "error" {
			return nil, &Error{Msg: "the workflow failed in ComfyUI: " + h.failure()}
		}
		if !h.Status.Completed {
			continue
		}
		return c.fetch(ctx, base, apiKey, h)
	}
}

type historyEntry struct {
	Status struct {
		StatusStr string              `json:"status_str"`
		Completed bool                `json:"completed"`
		Messages  [][]json.RawMessage `json:"messages"`
	} `json:"status"`
	Outputs map[string]struct {
		Images []struct {
			Filename  string `json:"filename"`
			Subfolder string `json:"subfolder"`
			Type      string `json:"type"`
		} `json:"images"`
	} `json:"outputs"`
}

// failure is what ComfyUI said went wrong, from its execution_error.
func (h historyEntry) failure() string {
	for _, m := range h.Status.Messages {
		if len(m) < 2 || string(m[0]) != `"execution_error"` {
			continue
		}
		var e struct {
			NodeType string `json:"node_type"`
			Message  string `json:"exception_message"`
		}
		if json.Unmarshal(m[1], &e) == nil && e.Message != "" {
			return strings.TrimSpace(e.NodeType + ": " + strings.TrimSpace(e.Message))
		}
	}
	return "no reason given"
}

func (c *Client) fetch(ctx context.Context, base, apiKey string, h historyEntry) ([]Image, error) {
	nodes := make([]string, 0, len(h.Outputs))
	for id := range h.Outputs {
		nodes = append(nodes, id)
	}
	// Numerically where they are numbers, which is how ComfyUI names them.
	slices.SortFunc(nodes, func(a, b string) int {
		x, ea := strconv.Atoi(a)
		y, eb := strconv.Atoi(b)
		if ea == nil && eb == nil {
			return x - y
		}
		return strings.Compare(a, b)
	})
	var out []Image
	for _, n := range nodes {
		for _, im := range h.Outputs[n].Images {
			q := url.Values{"filename": {im.Filename}, "subfolder": {im.Subfolder}, "type": {im.Type}}
			img, err := c.get(ctx, base+"/view?"+q.Encode(), apiKey)
			if err != nil {
				return nil, err
			}
			out = append(out, img)
		}
	}
	if len(out) == 0 {
		return nil, &Error{Msg: "the workflow finished without a picture: it has no SaveImage or PreviewImage node"}
	}
	return out, nil
}

// cancel takes a workflow out of the queue, or stops it if it is the one
// running. Best effort: the server may be gone, which also stops it.
func (c *Client) cancel(ctx context.Context, base, apiKey, id string) {
	ctx, done := context.WithTimeout(ctx, 5*time.Second)
	defer done()
	body, _ := json.Marshal(map[string]any{"delete": []string{id}})
	_, _ = c.do(ctx, http.MethodPost, base+"/queue", apiKey, body, nil)
	var q struct {
		Running [][]json.RawMessage `json:"queue_running"`
	}
	if _, err := c.do(ctx, http.MethodGet, base+"/queue", apiKey, nil, &q); err != nil {
		return
	}
	for _, r := range q.Running {
		if len(r) > 1 && string(r[1]) == strconv.Quote(id) {
			_, _ = c.do(ctx, http.MethodPost, base+"/interrupt", apiKey, []byte("{}"), nil)
			return
		}
	}
}

func (c *Client) client() *http.Client {
	if c.HTTP != nil {
		return c.HTTP
	}
	return http.DefaultClient
}

func (c *Client) request(ctx context.Context, method, u, apiKey string, body []byte) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, method, u, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	resp, err := c.client().Do(req)
	if err != nil {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		return nil, &Error{Msg: "ComfyUI could not be reached: " + err.Error(), Retryable: true}
	}
	return resp, nil
}

// do sends a request and reads a JSON answer into out. A 5xx is an error;
// any other status is returned with whatever body could be read.
func (c *Client) do(ctx context.Context, method, u, apiKey string, body []byte, out any) (int, error) {
	resp, err := c.request(ctx, method, u, apiKey, body)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
	if err != nil {
		return 0, &Error{Msg: "reading ComfyUI's answer: " + err.Error(), Retryable: true}
	}
	if resp.StatusCode >= 500 {
		return resp.StatusCode, &Error{Msg: fmt.Sprintf("ComfyUI answered %d: %s", resp.StatusCode, snippet(b)), Retryable: true}
	}
	if out != nil && len(b) > 0 {
		if err := json.Unmarshal(b, out); err != nil && resp.StatusCode == http.StatusOK {
			return resp.StatusCode, &Error{Msg: "ComfyUI's answer is not JSON: " + snippet(b)}
		}
	}
	return resp.StatusCode, nil
}

func (c *Client) get(ctx context.Context, u, apiKey string) (Image, error) {
	resp, err := c.request(ctx, http.MethodGet, u, apiKey, nil)
	if err != nil {
		return Image{}, err
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20))
	if err != nil {
		return Image{}, &Error{Msg: "reading a picture from ComfyUI: " + err.Error(), Retryable: true}
	}
	if resp.StatusCode != http.StatusOK {
		return Image{}, &Error{Msg: fmt.Sprintf("ComfyUI answered %d for a picture it made", resp.StatusCode)}
	}
	ct := resp.Header.Get("Content-Type")
	if ct == "" || ct == "application/octet-stream" {
		ct = http.DetectContentType(b)
	}
	if ct == "image/png" {
		b = withoutPNGText(b)
	}
	return Image{Data: b, ContentType: ct}, nil
}

// withoutPNGText is a PNG without its text chunks. ComfyUI writes the whole
// workflow into every picture it saves — the prompt, and the checkpoint and
// nodes behind it — and what a picture says about how it was made is the
// operator's to publish, not something to hand every client. A file that is
// not a well-formed PNG is returned as it came.
func withoutPNGText(b []byte) []byte {
	const sig = "\x89PNG\r\n\x1a\n"
	if len(b) < len(sig) || string(b[:len(sig)]) != sig {
		return b
	}
	out := make([]byte, 0, len(b))
	out = append(out, b[:len(sig)]...)
	for i := len(sig); i < len(b); {
		if i+12 > len(b) {
			return b
		}
		n := int(binary.BigEndian.Uint32(b[i:]))
		end := i + 12 + n
		if n < 0 || end > len(b) {
			return b
		}
		switch string(b[i+4 : i+8]) {
		case "tEXt", "iTXt", "zTXt":
		default:
			out = append(out, b[i:end]...)
		}
		i = end
	}
	return out
}

// refusal is what /prompt said when it would not queue a workflow.
func refusal(e, nodes json.RawMessage) string {
	var top struct {
		Message string `json:"message"`
		Details string `json:"details"`
	}
	msg := snippet(e)
	if json.Unmarshal(e, &top) == nil && top.Message != "" {
		msg = top.Message
		if top.Details != "" {
			msg += ": " + top.Details
		}
	}
	var perNode map[string]struct {
		ClassType string `json:"class_type"`
		Errors    []struct {
			Message string `json:"message"`
			Details string `json:"details"`
		} `json:"errors"`
	}
	if json.Unmarshal(nodes, &perNode) == nil {
		var parts []string
		for _, n := range perNode {
			for _, er := range n.Errors {
				parts = append(parts, strings.TrimSpace(n.ClassType+": "+er.Message+" "+er.Details))
			}
		}
		slices.Sort(parts)
		if len(parts) > 0 {
			msg += " (" + strings.Join(parts, "; ") + ")"
		}
	}
	if msg == "" {
		msg = "no reason given"
	}
	return msg
}

func snippet(b []byte) string {
	s := strings.TrimSpace(string(b))
	if len(s) > 300 {
		s = s[:300] + "…"
	}
	return s
}