package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

// A small client for ComfyUI's HTTP API (the engine runs in a container on 127.0.0.1, see wxch-runner).

type comfyOut struct{ Filename, Subfolder, Type string }

type comfyClient struct {
	base string
	hc   *http.Client
}

func newComfy(base string) *comfyClient {
	return &comfyClient{base: base, hc: &http.Client{Timeout: 5 * time.Minute}}
}

func (c *comfyClient) do(ctx context.Context, method, path string, body io.Reader, ctype string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, method, c.base+path, body)
	if err != nil {
		return nil, err
	}
	if ctype != "" {
		req.Header.Set("Content-Type", ctype)
	}
	return c.hc.Do(req)
}

// Upload puts a file into the engine's input folder under subfolder, and returns the name a workflow must use for it.
func (c *comfyClient) Upload(ctx context.Context, subfolder, name string, data io.Reader) (string, error) {
	var buf bytes.Buffer
	mw := multipart.NewWriter(&buf)
	mw.WriteField("subfolder", subfolder)
	mw.WriteField("type", "input")
	mw.WriteField("overwrite", "true")
	fw, err := mw.CreateFormFile("image", name)
	if err != nil {
		return "", err
	}
	if _, err := io.Copy(fw, data); err != nil {
		return "", err
	}
	mw.Close()
	resp, err := c.do(ctx, "POST", "/upload/image", &buf, mw.FormDataContentType())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 500))
		return "", fmt.Errorf("the engine refused the file (%d): %s", resp.StatusCode, strings.TrimSpace(string(b)))
	}
	var r struct{ Name, Subfolder string }
	json.NewDecoder(resp.Body).Decode(&r)
	if r.Name == "" {
		r.Name = name
	}
	if r.Subfolder != "" {
		return r.Subfolder + "/" + r.Name, nil
	}
	return r.Name, nil
}

// Run submits a workflow and waits for it to finish. onProgress gets 0..1 as the engine reports it. When ctx ends the
// engine is told to stop. It returns the files the workflow saved.
func (c *comfyClient) Run(ctx context.Context, wf map[string]any, onProgress func(float64)) ([]comfyOut, error) {
	cid := randID()
	pctx, cancelWS := context.WithCancel(ctx)
	defer cancelWS()
	go c.watchProgress(pctx, cid, onProgress)

	b, _ := json.Marshal(map[string]any{"prompt": wf, "client_id": cid})
	resp, err := c.do(ctx, "POST", "/prompt", bytes.NewReader(b), "application/json")
	if err != nil {
		return nil, err
	}
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("the engine rejected the job: %s", comfyError(raw))
	}
	var sub struct {
		PromptID string `json:"prompt_id"`
	}
	json.Unmarshal(raw, &sub)
	if sub.PromptID == "" {
		return nil, fmt.Errorf("the engine gave no job id")
	}
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			c.interrupt(sub.PromptID)
			return nil, ctx.Err()
		case <-tick.C:
		}
		hr, err := c.do(ctx, "GET", "/history/"+url.PathEscape(sub.PromptID), nil, "")
		if err != nil {
			if ctx.Err() != nil {
				continue
			}
			return nil, err
		}
		var h map[string]struct {
			Outputs map[string]map[string]json.RawMessage `json:"outputs"`
			Status  struct {
				StatusStr string            `json:"status_str"`
				Completed bool              `json:"completed"`
				Messages  []json.RawMessage `json:"messages"`
			} `json:"status"`
		}
		json.NewDecoder(hr.Body).Decode(&h)
		hr.Body.Close()
		e, ok := h[sub.PromptID]
		if !ok {
			continue
		}
		if e.Status.StatusStr == "error" {
			return nil, fmt.Errorf("the engine failed: %s", lastErrorMessage(e.Status.Messages))
		}
		if !e.Status.Completed && e.Status.StatusStr != "success" {
			continue
		}
		var outs []comfyOut
		for _, node := range e.Outputs {
			for _, key := range []string{"images", "gifs", "videos", "video", "audio"} {
				var list []struct{ Filename, Subfolder, Type string }
				if json.Unmarshal(node[key], &list) != nil {
					continue
				}
				for _, o := range list {
					if o.Type == "output" && o.Filename != "" {
						outs = append(outs, comfyOut{o.Filename, o.Subfolder, o.Type})
					}
				}
			}
		}
		c.forget(sub.PromptID)
		return outs, nil
	}
}

func lastErrorMessage(msgs []json.RawMessage) string {
	for i := len(msgs) - 1; i >= 0; i-- {
		var m []json.RawMessage
		if json.Unmarshal(msgs[i], &m) != nil || len(m) < 2 {
			continue
		}
		var kind string
		json.Unmarshal(m[0], &kind)
		if kind != "execution_error" {
			continue
		}
		var d struct {
			NodeType         string `json:"node_type"`
			ExceptionMessage string `json:"exception_message"`
		}
		json.Unmarshal(m[1], &d)
		return strings.TrimSpace(d.NodeType + ": " + d.ExceptionMessage)
	}
	return "see the engine log"
}

func comfyError(raw []byte) string {
	var e struct {
		Error      struct{ Message, Details string } `json:"error"`
		NodeErrors map[string]struct {
			Errors []struct{ Message, Details string } `json:"errors"`
		} `json:"node_errors"`
	}
	json.Unmarshal(raw, &e)
	msg := e.Error.Message
	for _, n := range e.NodeErrors {
		for _, x := range n.Errors {
			msg += "; " + x.Message + " " + x.Details
		}
	}
	if msg == "" {
		msg = strings.TrimSpace(string(raw))
		if len(msg) > 300 {
			msg = msg[:300]
		}
	}
	return msg
}

func (c *comfyClient) interrupt(pid string) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	b, _ := json.Marshal(map[string]any{"delete": []string{pid}})
	if r, err := c.do(ctx, "POST", "/queue", bytes.NewReader(b), "application/json"); err == nil {
		r.Body.Close()
	}
	if r, err := c.do(ctx, "POST", "/interrupt", nil, ""); err == nil {
		r.Body.Close()
	}
}

func (c *comfyClient) forget(pid string) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	b, _ := json.Marshal(map[string]any{"delete": []string{pid}})
	if r, err := c.do(ctx, "POST", "/history", bytes.NewReader(b), "application/json"); err == nil {
		r.Body.Close()
	}
}

// Fetch opens a saved file.
func (c *comfyClient) Fetch(ctx context.Context, o comfyOut) (io.ReadCloser, int64, error) {
	q := url.Values{"filename": {o.Filename}, "subfolder": {o.Subfolder}, "type": {o.Type}}
	resp, err := c.do(ctx, "GET", "/view?"+q.Encode(), nil, "")
	if err != nil {
		return nil, 0, err
	}
	if resp.StatusCode != 200 {
		resp.Body.Close()
		return nil, 0, fmt.Errorf("could not read %s from the engine (%d)", o.Filename, resp.StatusCode)
	}
	return resp.Body, resp.ContentLength, nil
}

// watchProgress follows the engine's websocket for step counts. It is best effort: without it the job still runs.
func (c *comfyClient) watchProgress(ctx context.Context, cid string, onProgress func(float64)) {
	if onProgress == nil {
		return
	}
	u := strings.Replace(c.base, "http", "ws", 1) + "/ws?clientId=" + cid
	ws, _, err := websocket.DefaultDialer.DialContext(ctx, u, nil)
	if err != nil {
		return
	}
	defer ws.Close()
	go func() { <-ctx.Done(); ws.Close() }()
	for {
		mt, data, err := ws.ReadMessage()
		if err != nil {
			return
		}
		if mt != websocket.TextMessage {
			continue
		}
		var m struct {
			Type string                       `json:"type"`
			Data struct{ Value, Max float64 } `json:"data"`
		}
		if json.Unmarshal(data, &m) == nil && m.Type == "progress" && m.Data.Max > 0 {
			onProgress(m.Data.Value / m.Data.Max)
		}
	}
}