package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

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

// The job API for programs (api.wxch.org, with an API key), next to the studio page:
//
//	GET  /v1/capabilities               what can be asked for and whether a machine offers it now
//	POST /v1/jobs                       JSON {capability, input, files:{field:{name,data}}} (data is base64) or multipart
//	GET  /v1/jobs, /v1/jobs/{id}        the key's own jobs, with progress and result links
//	DELETE /v1/jobs/{id}                cancel (queued or running) or delete
//	GET  /v1/jobs/{id}/files/{name}     a result file (Bearer key, or the signed link in the job)
//
// and wrappers in the OpenAI shape that wait for the result (a few minutes for video, seconds to a minute for the rest):
//
//	POST /v1/images/generations   POST /v1/images/edits   POST /v1/audio/transcriptions   POST /v1/audio/translations

const (
	maxActivePerKey = 5
	fileLinkTTL     = time.Hour
	syncWait        = 20 * time.Minute
)

type apiErr struct {
	Code int
	Msg  string
	Type string
}

func (e *apiErr) write(w http.ResponseWriter) { apiError(w, e.Code, e.Msg, e.Type) }

func studioOwner(o string) bool { return o == "" || o == "studio" }

// admitJob validates the input (the stored file names already in it), checks that a machine offers the capability and
// queues the job. On any failure the job's folder is removed.
func (s *Server) admitJob(owner string, cp modules.Capability, id string, input map[string]any) (*Job, *apiErr) {
	fail := func(code int, typ, msg string) (*Job, *apiErr) {
		os.RemoveAll(s.jobs.jobDir(id))
		return nil, &apiErr{code, msg, typ}
	}
	valid, err := cp.Validate(input)
	if err != nil {
		return fail(http.StatusBadRequest, "invalid_request", err.Error())
	}
	if o := s.hub.offers(s.jobs.busyCopy())[cp.ID]; o == nil || o.Online == 0 {
		return fail(http.StatusServiceUnavailable, "service_unavailable", "no machine offers this right now")
	}
	j, err := s.jobs.create(id, cp.ID, owner, valid)
	if err != nil {
		return fail(http.StatusTooManyRequests, "rate_limit_exceeded", err.Error())
	}
	return j, nil
}

type upload struct {
	Name string
	Data []byte
}

// enqueueBytes is for requests that arrive with their files in memory (JSON and the OpenAI forms).
func (s *Server) enqueueBytes(owner, capID string, input map[string]any, files map[string]upload) (*Job, *apiErr) {
	cp, ok := s.jobs.cat.Capability(capID)
	if !ok {
		return nil, &apiErr{http.StatusBadRequest, "unknown capability " + capID, "invalid_request"}
	}
	id := randHex(8)
	if input == nil {
		input = map[string]any{}
	}
	for field, up := range files {
		name := safeName(up.Name)
		if name == "" || !isFileField(cp, field) {
			os.RemoveAll(s.jobs.jobDir(id))
			return nil, &apiErr{http.StatusBadRequest, "unexpected file " + field, "invalid_request"}
		}
		max := cp.MaxInput
		if max <= 0 {
			max = 20 << 20
		}
		if len(up.Data) == 0 || int64(len(up.Data)) > max {
			os.RemoveAll(s.jobs.jobDir(id))
			return nil, &apiErr{http.StatusRequestEntityTooLarge, fmt.Sprintf("%s must be between 1 byte and %d MB", field, max>>20), "request_too_large"}
		}
		dir := filepath.Join(s.jobs.jobDir(id), "in")
		if err := os.MkdirAll(dir, 0750); err != nil {
			return nil, &apiErr{http.StatusInternalServerError, "storage error", "server_error"}
		}
		stored := field + "-" + name
		if err := os.WriteFile(filepath.Join(dir, stored), up.Data, 0640); err != nil {
			os.RemoveAll(s.jobs.jobDir(id))
			return nil, &apiErr{http.StatusInternalServerError, "storage error", "server_error"}
		}
		input[field] = stored
	}
	return s.admitJob(owner, cp, id, input)
}

// ---- links to result files ----

func (s *Server) linkKey() []byte {
	sum := sha256.Sum256([]byte("wxch-job-files|" + s.pwHash()))
	return sum[:]
}

func (s *Server) fileSig(id, name string, exp int64) string {
	m := hmac.New(sha256.New, s.linkKey())
	fmt.Fprintf(m, "%s|%s|%d", id, name, exp)
	return hex.EncodeToString(m.Sum(nil))[:32]
}

func (s *Server) fileURL(host, id, name string) string {
	exp := time.Now().Add(fileLinkTTL).Unix()
	return fmt.Sprintf("https://%s/v1/jobs/%s/files/%s?exp=%d&sig=%s", host, id, url.PathEscape(name), exp, s.fileSig(id, name, exp))
}

type apiOutput struct {
	Name string `json:"name"`
	Size int64  `json:"size"`
	Mime string `json:"mime"`
	URL  string `json:"url"` // valid for an hour; ask the job again for a fresh one
}

type apiJob struct {
	ID         string         `json:"id"`
	Object     string         `json:"object"`
	Capability string         `json:"capability"`
	Status     string         `json:"status"` // queued | running | done | failed | cancelled
	Progress   float64        `json:"progress"`
	Detail     string         `json:"detail,omitempty"`
	WaitingFor string         `json:"waiting_for,omitempty"` // while queued: no_machine | owner | busy
	Error      string         `json:"error,omitempty"`
	Outputs    []apiOutput    `json:"outputs,omitempty"`
	Text       string         `json:"text,omitempty"`
	Input      map[string]any `json:"input"`
	Created    int64          `json:"created"`
	Finished   int64          `json:"finished,omitempty"`
}

func (s *Server) apiView(j Job, host string) apiJob {
	v := apiJob{ID: j.ID, Object: "job", Capability: j.Capability, Status: j.Status, Progress: j.Pct, Detail: j.Detail, WaitingFor: j.Why,
		Error: j.Err, Text: j.Text, Input: map[string]any{}, Created: j.Created, Finished: j.Ended}
	for k, val := range j.Input { // the stored names of uploaded files mean nothing to the caller
		if _, isFile := val.(string); isFile && s.isFileInput(j.Capability, k) {
			continue
		}
		v.Input[k] = val
	}
	if j.Status == "done" {
		for _, o := range j.Outputs {
			v.Outputs = append(v.Outputs, apiOutput{o.Name, o.Size, o.Mime, s.fileURL(host, j.ID, o.Name)})
		}
	}
	return v
}

func (s *Server) isFileInput(capID, field string) bool {
	cp, ok := s.jobs.cat.Capability(capID)
	return ok && isFileField(cp, field)
}

// waitJob blocks until the job has ended, the caller went away (then the job is cancelled) or the time is up.
func (s *Server) waitJob(r *http.Request, id string) (Job, error) {
	deadline := time.NewTimer(syncWait)
	defer deadline.Stop()
	tick := time.NewTicker(500 * time.Millisecond)
	defer tick.Stop()
	for {
		j, ok := s.jobs.get(id)
		if !ok {
			return Job{}, fmt.Errorf("the job disappeared")
		}
		switch j.Status {
		case "done", "failed", "cancelled":
			return j, nil
		}
		select {
		case <-r.Context().Done():
			s.jobs.cancel(s.hub, id)
			return j, r.Context().Err()
		case <-deadline.C:
			s.jobs.cancel(s.hub, id)
			return j, fmt.Errorf("the job did not finish in %v", syncWait)
		case <-tick.C:
		}
	}
}

func (s *Server) readOutput(id, name string) ([]byte, error) {
	return os.ReadFile(filepath.Join(s.jobs.jobDir(id), "out", safeName(name)))
}

// ---- handlers (registered by gatewayMux) ----

type keyAuth func(w http.ResponseWriter, r *http.Request) (APIKey, bool)

func keyOwner(k APIKey) string { return "key:" + k.ID }

func (s *Server) registerJobAPI(mux *http.ServeMux, auth keyAuth, cors func(http.ResponseWriter)) {
	wrap := func(h func(w http.ResponseWriter, r *http.Request, key APIKey)) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			cors(w)
			if r.Method == http.MethodOptions {
				return
			}
			key, ok := auth(w, r)
			if !ok {
				return
			}
			h(w, r, key)
		}
	}
	mux.HandleFunc("/v1/capabilities", wrap(func(w http.ResponseWriter, r *http.Request, key APIKey) {
		writeJSON(w, map[string]any{"object": "list", "data": s.studioCaps()})
	}))
	mux.HandleFunc("/v1/jobs", wrap(func(w http.ResponseWriter, r *http.Request, key APIKey) {
		switch r.Method {
		case http.MethodGet:
			var data []apiJob
			for _, j := range s.jobs.list("key:" + key.ID) {
				data = append(data, s.apiView(j, r.Host))
			}
			writeJSON(w, map[string]any{"object": "list", "data": data})
		case http.MethodPost:
			s.apiSubmit(w, r, key)
		default:
			apiError(w, http.StatusMethodNotAllowed, "use GET or POST", "method_not_allowed")
		}
	}))
	// result files: the signed link in a job needs no key (an image tag cannot send one)
	mux.HandleFunc("/v1/jobs/", func(w http.ResponseWriter, r *http.Request) {
		cors(w)
		if r.Method == http.MethodOptions {
			return
		}
		parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/v1/jobs/"), "/")
		id := parts[0]
		if len(parts) == 3 && parts[1] == "files" && r.Method == http.MethodGet {
			s.apiFile(w, r, id, parts[2], auth)
			return
		}
		key, ok := auth(w, r)
		if !ok {
			return
		}
		j, found := s.jobs.get(id)
		if !found || j.Owner != keyOwner(key) {
			apiError(w, http.StatusNotFound, "no such job", "not_found")
			return
		}
		switch {
		case len(parts) == 1 && r.Method == http.MethodGet:
			writeJSON(w, s.apiView(j, r.Host))
		case len(parts) == 1 && r.Method == http.MethodDelete:
			s.jobs.remove(s.hub, id)
			writeJSON(w, map[string]any{"id": id, "object": "job", "deleted": true})
		default:
			apiError(w, http.StatusNotFound, "not found", "not_found")
		}
	})
	mux.HandleFunc("/v1/images/generations", wrap(s.openaiImages))
	mux.HandleFunc("/v1/images/edits", wrap(s.openaiImages))
	mux.HandleFunc("/v1/audio/transcriptions", wrap(s.openaiAudio))
	mux.HandleFunc("/v1/audio/translations", wrap(s.openaiAudio))
}

func (s *Server) apiFile(w http.ResponseWriter, r *http.Request, id, rawName string, auth keyAuth) {
	name := safeName(rawName)
	sig := r.URL.Query().Get("sig")
	var key APIKey
	if sig == "" { // without a signed link the caller must be the key that made the job
		var ok bool
		if key, ok = auth(w, r); !ok {
			return
		}
	}
	j, ok := s.jobs.get(id)
	if !ok || name == "" || j.Status != "done" {
		apiError(w, http.StatusNotFound, "no such file", "not_found")
		return
	}
	allowed := false
	if sig != "" {
		exp, _ := strconv.ParseInt(r.URL.Query().Get("exp"), 10, 64)
		allowed = exp > time.Now().Unix() && hmac.Equal([]byte(sig), []byte(s.fileSig(id, name, exp)))
	} else {
		allowed = j.Owner == keyOwner(key)
	}
	if !allowed {
		apiError(w, http.StatusNotFound, "no such file", "not_found")
		return
	}
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("Content-Type", mimeFor(name))
	http.ServeFile(w, r, filepath.Join(s.jobs.jobDir(id), "out", name))
}

// apiSubmit: POST /v1/jobs
func (s *Server) apiSubmit(w http.ResponseWriter, r *http.Request, key APIKey) {
	done, why := s.keys.admit(key.ID)
	if done == nil {
		w.Header().Set("Retry-After", "10")
		apiError(w, http.StatusTooManyRequests, why, "rate_limit_exceeded")
		return
	}
	defer done()
	var j *Job
	var e *apiErr
	if strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/") {
		j, e = s.submitMultipart(w, r, keyOwner(key))
	} else {
		var req struct {
			Capability string                                 `json:"capability"`
			Input      map[string]any                         `json:"input"`
			Files      map[string]struct{ Name, Data string } `json:"files"`
		}
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 96<<20)).Decode(&req) != nil || req.Capability == "" {
			apiError(w, http.StatusBadRequest, "body must be JSON with a capability", "invalid_json")
			return
		}
		files := map[string]upload{}
		for f, x := range req.Files {
			raw, err := base64.StdEncoding.DecodeString(x.Data)
			if err != nil {
				apiError(w, http.StatusBadRequest, "files."+f+".data must be base64", "invalid_request")
				return
			}
			files[f] = upload{x.Name, raw}
		}
		j, e = s.enqueueBytes(keyOwner(key), req.Capability, req.Input, files)
	}
	s.keys.record(key.ID, 0, 0, e != nil)
	if e != nil {
		e.write(w)
		return
	}
	w.WriteHeader(http.StatusCreated)
	writeJSON(w, s.apiView(*j, r.Host))
}

// finishSync waits for a job and answers the OpenAI-shaped caller; ok is false when an error was written.
func (s *Server) finishSync(w http.ResponseWriter, r *http.Request, key APIKey, j *Job) (Job, bool) {
	res, err := s.waitJob(r, j.ID)
	failed := err != nil || res.Status != "done"
	s.keys.record(key.ID, 0, 0, failed)
	switch {
	case err != nil:
		apiError(w, http.StatusGatewayTimeout, err.Error(), "timeout")
	case res.Status == "cancelled":
		apiError(w, http.StatusConflict, "the job was cancelled", "cancelled")
	case res.Status == "failed":
		apiError(w, http.StatusInternalServerError, res.Err, "job_failed")
	default:
		return res, true
	}
	return res, false
}

func formFile(r *http.Request, field string) (upload, error) {
	f, h, err := r.FormFile(field)
	if err != nil {
		return upload{}, err
	}
	defer f.Close()
	b, err := io.ReadAll(io.LimitReader(f, 128<<20))
	return upload{h.Filename, b}, err
}

// openaiImages: /v1/images/generations (JSON) and /v1/images/edits (multipart)
func (s *Server) openaiImages(w http.ResponseWriter, r *http.Request, key APIKey) {
	if r.Method != http.MethodPost {
		apiError(w, http.StatusMethodNotAllowed, "use POST", "method_not_allowed")
		return
	}
	done, why := s.keys.admit(key.ID)
	if done == nil {
		w.Header().Set("Retry-After", "10")
		apiError(w, http.StatusTooManyRequests, why, "rate_limit_exceeded")
		return
	}
	defer done()
	edit := strings.HasSuffix(r.URL.Path, "/edits")
	var prompt, size, format string
	n := 1
	files := map[string]upload{}
	if edit {
		if r.ParseMultipartForm(64<<20) != nil {
			apiError(w, http.StatusBadRequest, "send the form as multipart/form-data", "invalid_request")
			return
		}
		up, err := formFile(r, "image")
		if err != nil {
			apiError(w, http.StatusBadRequest, "image is required", "invalid_request")
			return
		}
		files["image"] = up
		prompt, format = r.FormValue("prompt"), r.FormValue("response_format")
	} else {
		var req struct {
			Prompt         string `json:"prompt"`
			N              int    `json:"n"`
			Size           string `json:"size"`
			ResponseFormat string `json:"response_format"`
		}
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&req) != nil {
			apiError(w, http.StatusBadRequest, "body must be JSON", "invalid_json")
			return
		}
		prompt, size, format = req.Prompt, req.Size, req.ResponseFormat
		if req.N > 0 {
			n = req.N
		}
	}
	input := map[string]any{"prompt": prompt}
	capID := "image.generate"
	if edit {
		capID = "image.edit"
	} else {
		if size != "" {
			input["size"] = size
		}
		input["count"] = n
	}
	j, e := s.enqueueBytes(keyOwner(key), capID, input, files)
	if e != nil {
		s.keys.record(key.ID, 0, 0, true)
		e.write(w)
		return
	}
	res, ok := s.finishSync(w, r, key, j)
	if !ok {
		return
	}
	var data []map[string]any
	for _, o := range res.Outputs {
		if !strings.HasPrefix(o.Mime, "image/") {
			continue
		}
		if format == "b64_json" {
			b, err := s.readOutput(res.ID, o.Name)
			if err != nil {
				continue
			}
			data = append(data, map[string]any{"b64_json": base64.StdEncoding.EncodeToString(b)})
		} else {
			data = append(data, map[string]any{"url": s.fileURL(r.Host, res.ID, o.Name)})
		}
	}
	writeJSON(w, map[string]any{"created": time.Now().Unix(), "data": data})
}

// openaiAudio: /v1/audio/transcriptions and /v1/audio/translations (multipart)
func (s *Server) openaiAudio(w http.ResponseWriter, r *http.Request, key APIKey) {
	if r.Method != http.MethodPost {
		apiError(w, http.StatusMethodNotAllowed, "use POST", "method_not_allowed")
		return
	}
	done, why := s.keys.admit(key.ID)
	if done == nil {
		w.Header().Set("Retry-After", "10")
		apiError(w, http.StatusTooManyRequests, why, "rate_limit_exceeded")
		return
	}
	defer done()
	if r.ParseMultipartForm(64<<20) != nil {
		apiError(w, http.StatusBadRequest, "send the form as multipart/form-data", "invalid_request")
		return
	}
	up, err := formFile(r, "file")
	if err != nil {
		apiError(w, http.StatusBadRequest, "file is required", "invalid_request")
		return
	}
	task := "transcribe"
	if strings.HasSuffix(r.URL.Path, "/translations") {
		task = "translate"
	}
	input := map[string]any{"task": task}
	if l := r.FormValue("language"); l != "" {
		input["language"] = l
	}
	j, e := s.enqueueBytes(keyOwner(key), "audio.transcribe", input, map[string]upload{"audio": up})
	if e != nil {
		s.keys.record(key.ID, 0, 0, true)
		e.write(w)
		return
	}
	res, ok := s.finishSync(w, r, key, j)
	if !ok {
		return
	}
	send := func(file, ctype string) {
		b, err := s.readOutput(res.ID, file)
		if err != nil {
			apiError(w, http.StatusInternalServerError, "the result is missing", "server_error")
			return
		}
		w.Header().Set("Content-Type", ctype)
		w.Write(b)
	}
	switch r.FormValue("response_format") {
	case "text":
		send("transcript.txt", "text/plain; charset=utf-8")
	case "srt":
		send("transcript.srt", "text/plain; charset=utf-8")
	case "vtt":
		send("transcript.vtt", "text/vtt; charset=utf-8")
	case "verbose_json":
		send("transcript.json", "application/json")
	default:
		writeJSON(w, map[string]string{"text": strings.TrimSpace(res.Text)})
	}
}