package main

import (
	"crypto/subtle"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"wxch.org/wxch/internal/modules"
)

// Jobs for the optional AI modules (image, video, speech). The server keeps the queue and the files; a machine that has
// the capability installed does the work (see cmd/wxch-agent/modules_jobs.go). One job runs per machine at a time. Files
// travel over HTTP (the machine fetches the input and uploads the result with its own token), only status goes over the
// websocket.

const (
	jobRetention    = 7 * 24 * time.Hour
	maxQueuedJobs   = 20
	maxOutputBytes  = 4 << 30
	maxJobTries     = 3
	jobRetryBackoff = 10 * time.Second
)

type JobOutput struct {
	Name string `json:"name"`
	Size int64  `json:"size"`
	Mime string `json:"mime"`
}

type Job struct {
	ID         string         `json:"id"`
	Capability string         `json:"capability"`
	Input      map[string]any `json:"input"`
	Status     string         `json:"status"` // queued | running | done | failed | cancelled
	Detail     string         `json:"detail,omitempty"`
	Pct        float64        `json:"pct"`
	Machine    string         `json:"machine,omitempty"`
	Err        string         `json:"err,omitempty"`
	Why        string         `json:"why,omitempty"`      // while queued: what it is waiting for
	WhyKind    string         `json:"why_kind,omitempty"` // when it waits for the owner: desktop | gpu | cooldown | paused
	RetryS     int            `json:"retry_s,omitempty"`  // and about how long, seconds
	Outputs    []JobOutput    `json:"outputs,omitempty"`
	Text       string         `json:"text,omitempty"`
	Created    int64          `json:"created"`
	Owner      string         `json:"owner,omitempty"` // "studio" (or empty) or "key:<id>" for a job made with an API key
	Seq        int64          `json:"seq"`             // order of creation: jobs made in the same second stay in order
	Started    int64          `json:"started,omitempty"`
	Ended      int64          `json:"ended,omitempty"`
	Tries      int            `json:"tries"`
	notBefore  time.Time
}

type JobStore struct {
	mu     sync.Mutex
	dir    string
	jobs   map[string]*Job
	busy   map[string]string // machine -> job id
	cat    *modules.Catalog
	wake   chan struct{}
	seq    int64
	warmAt map[string]time.Time // machine|capability -> when it was last asked to warm up
}

func newJobStore(dir string, cat *modules.Catalog) *JobStore {
	js := &JobStore{dir: dir, jobs: map[string]*Job{}, busy: map[string]string{}, cat: cat, wake: make(chan struct{}, 1)}
	os.MkdirAll(dir, 0750)
	if b, err := os.ReadFile(filepath.Join(dir, "jobs.json")); err == nil {
		var list []*Job
		if json.Unmarshal(b, &list) == nil {
			for _, j := range list {
				if j.Status == "running" { // the server restarted while it ran: the machine will report nothing more
					j.Status, j.Machine, j.Detail = "queued", "", ""
				}
				js.jobs[j.ID] = j
				if j.Seq > js.seq {
					js.seq = j.Seq
				}
			}
		}
	}
	return js
}

func (js *JobStore) saveLocked() {
	list := make([]*Job, 0, len(js.jobs))
	for _, j := range js.jobs {
		list = append(list, j)
	}
	sort.Slice(list, func(i, k int) bool { return list[i].Seq < list[k].Seq })
	b, _ := json.Marshal(list)
	p := filepath.Join(js.dir, "jobs.json")
	if os.WriteFile(p+".tmp", b, 0640) == nil {
		os.Rename(p+".tmp", p)
	}
}

func (js *JobStore) kick() {
	select {
	case js.wake <- struct{}{}:
	default:
	}
}

func (js *JobStore) jobDir(id string) string { return filepath.Join(js.dir, id) }

// safeName keeps only a plain file name.
func safeName(n string) string {
	n = filepath.Base(strings.ReplaceAll(n, "\\", "/"))
	if n == "." || n == ".." || n == "/" || n == "" || len(n) > 200 || strings.ContainsAny(n, "\x00") {
		return ""
	}
	return n
}

// offers lists, per capability, how many online machines have it installed and whether one of them is free now.
type offer struct {
	Online int    // machines with it installed that are online
	Free   int    // of those, machines that take work now
	Why    string // when none is free: owner | busy | ""
	Kind   string // when the owner: desktop | gpu | cooldown | paused (what the machine reports, no program names)
	RetryS int    // the shortest wait among those machines, seconds (0 = unknown)
}

func (h *Hub) offers(busy map[string]string) map[string]*offer {
	h.mu.Lock()
	defer h.mu.Unlock()
	out := map[string]*offer{}
	for n := range h.agents {
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second {
			continue
		}
		for id := range b.modules {
			o := out[id]
			if o == nil {
				o = &offer{}
				out[id] = o
			}
			o.Online++
			switch {
			case b.mode != "" && b.mode != "available":
				if o.Why == "" {
					o.Why = "owner"
				}
				if w := b.waitS(); o.Kind == "" || (w > 0 && (o.RetryS == 0 || w < o.RetryS)) {
					o.Kind, o.RetryS = b.kind, w // the machine that frees up first
				}
			case busy[n] != "":
				if o.Why == "" {
					o.Why = "busy"
				}
			default:
				o.Free++
			}
		}
	}
	return out
}

// pick chooses the free machine with the most memory left.
func (js *JobStore) pick(h *Hub, capID string) (string, *conn) {
	h.mu.Lock()
	defer h.mu.Unlock()
	best, bestMem := "", -1.0
	for n := range h.agents {
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second || !b.modules[capID] || js.busy[n] != "" {
			continue
		}
		if b.mode != "" && b.mode != "available" {
			continue
		}
		if b.memAvail > bestMem {
			best, bestMem = n, b.memAvail
		}
	}
	if best == "" {
		return "", nil
	}
	return best, h.agents[best]
}

// schedule hands queued jobs (oldest first) to free machines.
func (js *JobStore) schedule(h *Hub) {
	js.mu.Lock()
	var queued []*Job
	for _, j := range js.jobs {
		if j.Status == "queued" && time.Now().After(j.notBefore) {
			queued = append(queued, j)
		}
	}
	sort.Slice(queued, func(i, k int) bool { return queued[i].Seq < queued[k].Seq })
	js.mu.Unlock()
	offers := h.offers(js.busyCopy())
	for _, j := range queued {
		js.mu.Lock()
		if j.Status != "queued" {
			js.mu.Unlock()
			continue
		}
		machine, c := js.pick(h, j.Capability)
		if c == nil {
			switch o := offers[j.Capability]; {
			case o == nil || o.Online == 0:
				j.Why = "no_machine"
			case o.Why == "owner":
				j.Why, j.WhyKind, j.RetryS = "owner", o.Kind, o.RetryS
			default:
				j.Why = "busy"
			}
			js.mu.Unlock()
			continue
		}
		j.Status, j.Machine, j.Why, j.WhyKind, j.RetryS, j.Started, j.Detail, j.Pct = "running", machine, "", "", 0, time.Now().Unix(), "starting", 0
		j.Tries++
		js.busy[machine] = j.ID
		params, _ := json.Marshal(map[string]any{"input": j.Input})
		id, capID := j.ID, j.Capability
		js.saveLocked()
		js.mu.Unlock()
		c.send(Msg{T: "job", ID: id, Method: capID, Params: params})
	}
}

func (js *JobStore) busyCopy() map[string]string {
	js.mu.Lock()
	defer js.mu.Unlock()
	m := map[string]string{}
	for k, v := range js.busy {
		m[k] = v
	}
	return m
}

func (js *JobStore) run(h *Hub) {
	t := time.NewTicker(2 * time.Second)
	cleanup := time.NewTicker(time.Hour)
	for {
		select {
		case <-t.C:
		case <-js.wake:
		case <-cleanup.C:
			js.expire()
		}
		js.schedule(h)
	}
}

// expire deletes finished jobs and their files after the retention time.
func (js *JobStore) expire() {
	js.mu.Lock()
	defer js.mu.Unlock()
	cut := time.Now().Add(-jobRetention).Unix()
	for id, j := range js.jobs {
		if j.Status != "queued" && j.Status != "running" && j.Ended > 0 && j.Ended < cut {
			os.RemoveAll(js.jobDir(id))
			delete(js.jobs, id)
		}
	}
	js.saveLocked()
}

// onAgentMsg takes job_progress and job_done from a machine.
func (js *JobStore) onAgentMsg(machine string, m Msg) {
	js.mu.Lock()
	defer js.mu.Unlock()
	j := js.jobs[m.ID]
	if j == nil || j.Machine != machine || j.Status != "running" {
		return
	}
	switch m.T {
	case "job_progress":
		var p struct {
			Status string
			Pct    float64
		}
		json.Unmarshal(m.Data, &p)
		j.Detail, j.Pct = p.Status, p.Pct
	case "job_done":
		delete(js.busy, machine)
		j.Ended = time.Now().Unix()
		switch {
		case m.OK:
			var r struct {
				Outputs []JobOutput
				Text    string
			}
			json.Unmarshal(m.Data, &r)
			// what the machine says it stored must really be there
			for _, o := range r.Outputs {
				if fi, err := os.Stat(filepath.Join(js.jobDir(j.ID), "out", safeName(o.Name))); err == nil {
					o.Size = fi.Size()
					j.Outputs = append(j.Outputs, o)
				}
			}
			j.Status, j.Text, j.Detail, j.Pct = "done", r.Text, "", 1
			if len(j.Outputs) == 0 {
				j.Status, j.Err = "failed", "the machine finished but no result arrived"
			}
		case m.Status == 499:
			j.Status, j.Detail = "cancelled", ""
		case m.Status == 409 && j.Tries < 50: // the machine could not take it now (owner back, busy): try again later
			js.requeueLocked(j, m.Err)
		default:
			j.Status, j.Err, j.Detail = "failed", m.Err, ""
		}
		js.saveLocked()
		js.kick()
		return
	}
}

func (js *JobStore) requeueLocked(j *Job, why string) {
	j.Status, j.Machine, j.Detail, j.Pct, j.Ended = "queued", "", "", 0, 0
	j.notBefore = time.Now().Add(jobRetryBackoff)
	if j.Tries > 0 && strings.Contains(why, "busy with another job") {
		j.Tries-- // not the job's fault
	}
}

// machineGone requeues what the machine was doing.
func (js *JobStore) machineGone(machine string) {
	js.mu.Lock()
	defer js.mu.Unlock()
	id := js.busy[machine]
	delete(js.busy, machine)
	if j := js.jobs[id]; j != nil && j.Status == "running" {
		if j.Tries >= maxJobTries {
			j.Status, j.Err, j.Ended = "failed", "the machine disconnected", time.Now().Unix()
		} else {
			os.RemoveAll(filepath.Join(js.jobDir(id), "out")) // a half-sent result
			js.requeueLocked(j, "machine disconnected")
		}
		js.saveLocked()
		js.kick()
	}
}

// ---- warming ----

const warmEvery = 45 * time.Second

// warm asks a machine that offers the capability to start its engine ahead of time. It picks an online machine that takes
// work (a machine the owner is using is left alone), and asks each machine at most once per warmEvery for one capability.
func (s *Server) warm(capID string) bool {
	h := s.hub
	h.mu.Lock()
	var name string
	var c *conn
	for n, a := range h.agents {
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second || !b.modules[capID] || (b.mode != "" && b.mode != "available") {
			continue
		}
		if name == "" || n < name {
			name, c = n, a
		}
	}
	h.mu.Unlock()
	if c == nil {
		return false
	}
	js := s.jobs
	js.mu.Lock()
	if js.warmAt == nil {
		js.warmAt = map[string]time.Time{}
	}
	key := name + "|" + capID
	if time.Since(js.warmAt[key]) < warmEvery {
		js.mu.Unlock()
		return true // asked a moment ago
	}
	js.warmAt[key] = time.Now()
	js.mu.Unlock()
	c.send(Msg{T: "warm", Method: capID})
	return true
}

// ---- creating, listing, cancelling ----

func (js *JobStore) queuedCount() int {
	n := 0
	for _, j := range js.jobs {
		if j.Status == "queued" || j.Status == "running" {
			n++
		}
	}
	return n
}

// create adds a job whose validated input is already in place (input files stored under in/).
func (js *JobStore) create(id, capID, owner string, input map[string]any) (*Job, error) {
	js.mu.Lock()
	defer js.mu.Unlock()
	if js.queuedCount() >= maxQueuedJobs {
		return nil, fmt.Errorf("too many jobs are waiting; try again when some have finished")
	}
	if !studioOwner(owner) {
		active := 0
		for _, j := range js.jobs {
			if j.Owner == owner && (j.Status == "queued" || j.Status == "running") {
				active++
			}
		}
		if active >= maxActivePerKey {
			return nil, fmt.Errorf("at most %d jobs of one key may wait or run at the same time", maxActivePerKey)
		}
	}
	js.seq++
	if studioOwner(owner) {
		owner = "studio"
	}
	j := &Job{ID: id, Capability: capID, Owner: owner, Input: input, Status: "queued", Created: time.Now().Unix(), Seq: js.seq}
	js.jobs[id] = j
	js.saveLocked()
	js.kick()
	return j, nil
}

func (js *JobStore) get(id string) (Job, bool) {
	js.mu.Lock()
	defer js.mu.Unlock()
	j := js.jobs[id]
	if j == nil {
		return Job{}, false
	}
	return *j, true
}

// list returns the newest jobs of one owner ("studio" or "key:<id>").
func (js *JobStore) list(owner string) []Job {
	js.mu.Lock()
	defer js.mu.Unlock()
	out := make([]Job, 0, len(js.jobs))
	for _, j := range js.jobs {
		if j.Owner == owner || (owner == "studio" && studioOwner(j.Owner)) {
			out = append(out, *j)
		}
	}
	sort.Slice(out, func(i, k int) bool { return out[i].Seq > out[k].Seq })
	if len(out) > 60 {
		out = out[:60]
	}
	return out
}

// cancel stops a queued job at once and asks the machine to stop a running one.
func (js *JobStore) cancel(h *Hub, id string) bool {
	js.mu.Lock()
	j := js.jobs[id]
	if j == nil {
		js.mu.Unlock()
		return false
	}
	var c *conn
	switch j.Status {
	case "queued":
		j.Status, j.Ended = "cancelled", time.Now().Unix()
	case "running":
		h.mu.Lock()
		c = h.agents[j.Machine]
		h.mu.Unlock()
		j.Detail = "cancelling"
	}
	js.saveLocked()
	js.mu.Unlock()
	if c != nil {
		c.send(Msg{T: "job_cancel", ID: id})
	}
	return true
}

func (js *JobStore) remove(h *Hub, id string) bool {
	js.mu.Lock()
	j := js.jobs[id]
	if j == nil {
		js.mu.Unlock()
		return false
	}
	running := j.Status == "running"
	js.mu.Unlock()
	if running {
		js.cancel(h, id)
		return true
	}
	js.mu.Lock()
	delete(js.jobs, id)
	js.saveLocked()
	js.mu.Unlock()
	os.RemoveAll(js.jobDir(id))
	return true
}

// ---- file transfer with the machines ----

// agentName finds which machine a token belongs to.
func (s *Server) agentName(r *http.Request) string {
	tok := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
	if tok == "" {
		return ""
	}
	for n, t := range s.agentTokens() {
		if subtle.ConstantTimeCompare([]byte(t), []byte(tok)) == 1 {
			return n
		}
	}
	return ""
}

// agentJobFile: GET dir=in sends an input file to the machine that runs the job; POST dir=out stores a result file.
func (s *Server) agentJobFile(w http.ResponseWriter, r *http.Request) {
	machine := s.agentName(r)
	if machine == "" {
		http.Error(w, "unauthorized", 401)
		return
	}
	q := r.URL.Query()
	id, name, dir := q.Get("job"), safeName(q.Get("name")), q.Get("dir")
	js := s.jobs
	j, ok := js.get(id)
	if !ok || j.Machine != machine || j.Status != "running" || name == "" || (dir != "in" && dir != "out") {
		http.Error(w, "no such job", 404)
		return
	}
	p := filepath.Join(js.jobDir(id), dir, name)
	if r.Method == http.MethodGet && dir == "in" {
		http.ServeFile(w, r, p)
		return
	}
	if r.Method != http.MethodPost || dir != "out" {
		http.Error(w, "bad request", 400)
		return
	}
	if err := os.MkdirAll(filepath.Dir(p), 0750); err != nil {
		http.Error(w, "storage error", 500)
		return
	}
	f, err := os.Create(p + ".part")
	if err != nil {
		http.Error(w, "storage error", 500)
		return
	}
	n, err := io.Copy(f, io.LimitReader(r.Body, maxOutputBytes+1))
	f.Close()
	if err != nil || n > maxOutputBytes {
		os.Remove(p + ".part")
		http.Error(w, "upload failed", 400)
		return
	}
	if err := os.Rename(p+".part", p); err != nil {
		http.Error(w, "storage error", 500)
		return
	}
	w.Write([]byte("ok"))
	log.Printf("job %s: %s stored (%d bytes) from %s", id, name, n, machine)
}