package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Download a model into ComfyUI's models folder on this machine. ComfyUI's own "Download" buttons make the browser fetch
// the file, which puts it on the computer of whoever clicked; the page we serve (comfy.) sends those clicks here instead,
// through the tunnel, and the machine downloads the file itself. Only the administrator reaches this (the server checks the
// login before anything goes through the tunnel), and only Hugging Face files that end up in a models sub-folder are taken.
//
//	POST /_wxch/fetch   {"url": "...", "dir": "vae", "name": "ae.safetensors"}  -> {"ok": true, "job": {...}}
//	GET  /_wxch/fetch                                                            -> {"jobs": [...]}

const fetchEngine = "comfyui"

var (
	fetchDirRe  = regexp.MustCompile(`^[a-z][a-z0-9_]{1,40}$`)
	fetchNameRe = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$`)
	fetchExts   = map[string]bool{".safetensors": true, ".ckpt": true, ".pt": true, ".pth": true, ".bin": true, ".gguf": true, ".sft": true, ".onnx": true}
)

type fetchJob struct {
	Dir     string `json:"dir"`
	Name    string `json:"name"`
	State   string `json:"state"` // starting, downloading, done, failed
	Total   int64  `json:"total"`
	Done    int64  `json:"done"`
	Err     string `json:"err,omitempty"`
	Started int64  `json:"started"`
	Ended   int64  `json:"ended,omitempty"`
}

var fetches = struct {
	sync.Mutex
	jobs map[string]*fetchJob
}{jobs: map[string]*fetchJob{}}

// fetchSizeFn finds the size of the file to download (a variable for tests).
var fetchSizeFn = fetchSize

// fetchDownload does the download (a variable for tests).
var fetchDownload = func(ctx context.Context, url, dest string, size int64, onBytes func(int64)) error {
	return downloadFile(ctx, url, dest, size, "", onBytes) // a file someone picked in ComfyUI: no published hash to check
}

// fetchHostOK says whether a host is Hugging Face (or one of its subdomains).
func fetchHostOK(h string) bool {
	h = strings.ToLower(h)
	for _, d := range []string{"huggingface.co", "hf.co"} {
		if h == d || strings.HasSuffix(h, "."+d) {
			return true
		}
	}
	return false
}

// cleanFetchURL accepts only an https Hugging Face file address and returns it without query or fragment; a "blob" page
// address is turned into the "resolve" address that serves the file.
func cleanFetchURL(raw string) (string, string, error) {
	u, err := url.Parse(strings.TrimSpace(raw))
	if err != nil || u.Scheme != "https" || u.User != nil || u.Port() != "" || !fetchHostOK(u.Hostname()) {
		return "", "", fmt.Errorf("only https://huggingface.co/... file addresses can be downloaded here")
	}
	p := strings.Replace(u.EscapedPath(), "/blob/", "/resolve/", 1)
	if !strings.Contains(p, "/resolve/") {
		return "", "", fmt.Errorf("this is not the address of a file on Hugging Face (it should contain /resolve/)")
	}
	u.RawQuery, u.Fragment, u.RawFragment = "", "", ""
	u.Path, _ = url.PathUnescape(p)
	u.RawPath = p
	return u.String(), path.Base(u.Path), nil
}

// fetchSize asks how big the file is: the first byte is requested, and the answer tells the whole size.
func fetchSize(ctx context.Context, raw string) (int64, error) {
	req, _ := http.NewRequestWithContext(ctx, "GET", raw, nil)
	if t := hfToken(); t != "" {
		req.Header.Set("Authorization", "Bearer "+t)
	}
	req.Header.Set("Range", "bytes=0-0")
	resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	io.Copy(io.Discard, io.LimitReader(resp.Body, 16))
	switch resp.StatusCode {
	case 206:
		if _, total, ok := strings.Cut(resp.Header.Get("Content-Range"), "/"); ok {
			if n, err := strconv.ParseInt(total, 10, 64); err == nil && n > 0 {
				return n, nil
			}
		}
	case 200:
		if resp.ContentLength > 0 {
			return resp.ContentLength, nil
		}
	case 401, 403:
		return 0, fmt.Errorf("Hugging Face refused access (a gated model needs a token in /etc/wxch/hf_token)")
	case 404:
		return 0, fmt.Errorf("the file was not found on Hugging Face")
	default:
		return 0, fmt.Errorf("Hugging Face answered HTTP %d", resp.StatusCode)
	}
	return 0, fmt.Errorf("could not find out how big the file is")
}

func fetchActive() int {
	n := 0
	for _, j := range fetches.jobs {
		if j.State == "starting" || j.State == "downloading" {
			n++
		}
	}
	return n
}

type fetchReq struct {
	URL  string `json:"url"`
	Dir  string `json:"dir"`
	Name string `json:"name"`
}

// startFetch checks a request and starts the download in the background.
func startFetch(cfg nodeConfig, req fetchReq) (fetchJob, error) {
	cat, err := loadCatalog()
	if err != nil {
		return fetchJob{}, err
	}
	e, ok := cat.Engine(fetchEngine)
	if !ok || len(capabilitiesUsing(cat, readModuleState(cfg), fetchEngine)) == 0 {
		return fetchJob{}, fmt.Errorf("no image or video module is installed on this machine")
	}
	clean, base, err := cleanFetchURL(req.URL)
	if err != nil {
		return fetchJob{}, err
	}
	name := strings.TrimSpace(req.Name)
	if name == "" {
		name = base
	}
	if !fetchNameRe.MatchString(name) || strings.Contains(name, "..") || !fetchExts[strings.ToLower(filepath.Ext(name))] {
		return fetchJob{}, fmt.Errorf("%q is not a model file name (.safetensors, .ckpt, .pt, .pth, .bin, .gguf)", name)
	}
	if !fetchDirRe.MatchString(req.Dir) {
		return fetchJob{}, fmt.Errorf("choose the models folder the file belongs in (vae, text_encoders, diffusion_models, loras, checkpoints, ...)")
	}
	root := filepath.Join(engineDir(cfg, e), filepath.FromSlash(e.ModelsRoot))
	if _, err := os.Stat(root); err != nil {
		return fetchJob{}, fmt.Errorf("the ComfyUI models folder is not there yet")
	}
	dest := filepath.Join(root, req.Dir, name)
	key := req.Dir + "/" + name
	fetches.Lock()
	if j := fetches.jobs[key]; j != nil && (j.State == "starting" || j.State == "downloading") {
		snap := *j
		fetches.Unlock()
		return snap, nil
	}
	if fetchActive() >= 3 {
		fetches.Unlock()
		return fetchJob{}, fmt.Errorf("three downloads are running; wait for one to finish")
	}
	j := &fetchJob{Dir: req.Dir, Name: name, State: "starting", Started: time.Now().Unix()}
	fetches.jobs[key] = j
	snap := *j
	fetches.Unlock()
	set := func(f func()) { fetches.Lock(); f(); fetches.Unlock() }
	fail := func(err error) { set(func() { j.State, j.Err, j.Ended = "failed", err.Error(), time.Now().Unix() }) }
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 6*time.Hour)
		defer cancel()
		if fi, err := os.Stat(dest); err == nil && fi.Size() > 0 {
			set(func() { j.State, j.Total, j.Done, j.Ended = "done", fi.Size(), fi.Size(), time.Now().Unix() })
			return
		}
		size, err := fetchSizeFn(ctx, clean)
		if err != nil {
			fail(err)
			return
		}
		if free, _, err := diskFree(modulesDir(cfg)); err == nil && uint64(size)+uint64(cfg.MinFreeGB)<<30 > free {
			fail(fmt.Errorf("not enough free disk space: the file is %.1f GB (+%d GB spare), %.1f GB is free", float64(size)/1e9, cfg.MinFreeGB, float64(free)/1e9))
			return
		}
		set(func() { j.State, j.Total = "downloading", size })
		var done int64
		err = fetchDownload(ctx, clean, dest, size, func(n int64) { done += n; set(func() { j.Done = done }) })
		if err != nil {
			fail(err)
			return
		}
		set(func() { j.State, j.Done, j.Ended = "done", size, time.Now().Unix() })
		logf("model downloaded into ComfyUI: %s (%.1f GB)", key, float64(size)/1e9)
	}()
	return snap, nil
}

func fetchList() []fetchJob {
	fetches.Lock()
	defer fetches.Unlock()
	var out []fetchJob
	for _, j := range fetches.jobs {
		out = append(out, *j)
	}
	sort.Slice(out, func(a, b int) bool { return out[a].Started > out[b].Started })
	if len(out) > 20 {
		out = out[:20]
	}
	return out
}

// serveLocalTun answers a tunnel request that is for the agent itself (paths under /_wxch/), not for the engine.
func serveLocalTun(ctx context.Context, m Msg, t *agentTun, method string, send func(Msg)) {
	reply := func(code int, v any) {
		b, _ := json.Marshal(v)
		h, _ := json.Marshal(map[string]string{"Content-Type": "application/json", "Cache-Control": "no-store"})
		send(Msg{T: "tun_res", ID: m.ID, Status: code, Params: h})
		sendChunks(send, m.ID, b)
		send(Msg{T: "tun_end", ID: m.ID})
	}
	p := m.D
	if i := strings.IndexByte(p, '?'); i >= 0 {
		p = p[:i]
	}
	if p != "/_wxch/fetch" {
		reply(404, map[string]string{"error": "not found"})
		return
	}
	switch method {
	case "GET":
		reply(200, map[string]any{"jobs": fetchList()})
	case "POST":
		var body []byte
	read:
		for len(body) <= 16<<10 {
			select {
			case <-ctx.Done():
				return
			case b := <-t.in:
				if b.T != "tun_body" {
					continue
				}
				if raw, err := base64.StdEncoding.DecodeString(b.D); err == nil {
					body = append(body, raw...)
				}
				if b.Status == 1 {
					break read
				}
			case <-time.After(15 * time.Second):
				reply(408, map[string]string{"error": "no body"})
				return
			}
		}
		var req fetchReq
		if json.Unmarshal(body, &req) != nil {
			reply(400, map[string]string{"error": "bad request"})
			return
		}
		j, err := startFetch(loadNodeConfig(), req)
		if err != nil {
			reply(400, map[string]string{"error": err.Error()})
			return
		}
		reply(200, map[string]any{"ok": true, "job": j})
	default:
		reply(405, map[string]string{"error": "method not allowed"})
	}
}