package main

import (
	"encoding/json"
	"io"
	"io/fs"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/bcrypt"

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

// studio.wxch.org: make images, video and transcripts with the AI modules that the machines have installed. Login is the
// web password (like the chat). A capability is listed as available only when an online machine has it installed: a
// provider that does not install a module simply does not offer it.

func (s *Server) studioAuthed(r *http.Request) bool {
	c, err := r.Cookie("wxch_studio")
	if err != nil {
		return false
	}
	return s.studioSess.valid(c.Value)
}

type studioCap struct {
	ID        string          `json:"id"`
	Title     modules.Text    `json:"title"`
	Desc      modules.Text    `json:"desc"`
	Note      modules.Text    `json:"note"`
	Inputs    []modules.Field `json:"inputs"`
	Outputs   []string        `json:"outputs"`
	MaxInput  int64           `json:"max_input_bytes"`
	MemBytes  int64           `json:"mem_bytes"` // memory a job needs on the machine (an estimate until measured)
	Measured  bool            `json:"measured"`
	Available bool            `json:"available"` // an online machine can take a job now
	Offered   bool            `json:"offered"`   // some online machine has it installed
	Why       string          `json:"why,omitempty"`
	WhyKind   string          `json:"why_kind,omitempty"` // owner: desktop | gpu | cooldown | paused
	RetryS    int             `json:"retry_s,omitempty"`  // about how long until a machine is free, seconds
}

func (s *Server) studioCaps() []studioCap {
	offers := s.hub.offers(s.jobs.busyCopy())
	var out []studioCap
	for _, c := range s.jobs.cat.Capabilities {
		sc := studioCap{ID: c.ID, Title: c.Title, Desc: c.Desc, Note: c.Note, Inputs: c.Inputs, Outputs: c.Outputs, MaxInput: c.MaxInput, MemBytes: c.MemBytes, Measured: c.Measured}
		if o := offers[c.ID]; o != nil && o.Online > 0 {
			sc.Offered, sc.Available, sc.Why = true, o.Free > 0, o.Why
			if !sc.Available && o.Why == "owner" {
				sc.WhyKind, sc.RetryS = o.Kind, o.RetryS
			}
		}
		out = append(out, sc)
	}
	return out
}

func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	json.NewEncoder(w).Encode(v)
}

func (s *Server) studioMux(sub fs.FS) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		b, _ := fs.ReadFile(sub, "studio.html")
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Cache-Control", "no-store")
		w.Write(b)
	})
	mux.HandleFunc("/studio/login", func(w http.ResponseWriter, r *http.Request) {
		ip := clientIP(r)
		if s.locked(ip) {
			http.Error(w, "too many attempts", http.StatusTooManyRequests)
			return
		}
		var req struct{ Password string }
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096)).Decode(&req) != nil ||
			bcrypt.CompareHashAndPassword([]byte(s.pwHash()), []byte(req.Password)) != nil {
			s.fail(ip)
			log.Printf("studio login failed from %s", ip)
			http.Error(w, "invalid", http.StatusUnauthorized)
			return
		}
		tok := randHex(32)
		s.studioSess.add(tok, 30*24*time.Hour)
		http.SetCookie(w, &http.Cookie{Name: "wxch_studio", Value: tok, Path: "/", HttpOnly: true,
			Secure: true, SameSite: http.SameSiteStrictMode, MaxAge: 30 * 24 * 3600})
		w.Write([]byte("ok"))
	})
	guard := func(h http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			if !s.studioAuthed(r) {
				http.Error(w, "unauthorized", http.StatusUnauthorized)
				return
			}
			h(w, r)
		}
	}
	mux.HandleFunc("/studio/api/capabilities", guard(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, map[string]any{"capabilities": s.studioCaps()})
	}))
	// the page asks for the engine behind the tab the person is looking at to be started, so it is ready when they press Start
	mux.HandleFunc("/studio/api/warm", guard(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "use POST", http.StatusMethodNotAllowed)
			return
		}
		var req struct{ Capability string }
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 1024)).Decode(&req) != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		if _, ok := s.jobs.cat.Capability(req.Capability); !ok {
			http.Error(w, "unknown capability", http.StatusBadRequest)
			return
		}
		writeJSON(w, map[string]any{"warming": s.warm(req.Capability)})
	}))
	mux.HandleFunc("/studio/api/jobs", guard(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			writeJSON(w, map[string]any{"jobs": s.jobs.list("studio")})
		case http.MethodPost:
			j, e := s.submitMultipart(w, r, "studio")
			if e != nil {
				http.Error(w, e.Msg, e.Code)
				return
			}
			writeJSON(w, j)
		default:
			http.Error(w, "use GET or POST", http.StatusMethodNotAllowed)
		}
	}))
	mux.HandleFunc("/studio/api/jobs/", guard(func(w http.ResponseWriter, r *http.Request) {
		rest := strings.TrimPrefix(r.URL.Path, "/studio/api/jobs/")
		parts := strings.Split(rest, "/")
		id := parts[0]
		j, ok := s.jobs.get(id)
		if !ok || !studioOwner(j.Owner) { // jobs made with an API key belong to that key
			http.NotFound(w, r)
			return
		}
		switch {
		case len(parts) == 1 && r.Method == http.MethodGet:
			writeJSON(w, j)
		case len(parts) == 1 && r.Method == http.MethodDelete:
			s.jobs.remove(s.hub, id)
			w.Write([]byte("ok"))
		case len(parts) == 2 && parts[1] == "cancel" && r.Method == http.MethodPost:
			s.jobs.cancel(s.hub, id)
			w.Write([]byte("ok"))
		case len(parts) == 3 && (parts[1] == "out" || parts[1] == "in") && r.Method == http.MethodGet:
			name := safeName(parts[2])
			if name == "" {
				http.NotFound(w, r)
				return
			}
			p := filepath.Join(s.jobs.jobDir(id), parts[1], name)
			w.Header().Set("X-Content-Type-Options", "nosniff")
			w.Header().Set("Content-Type", mimeFor(name))
			if r.URL.Query().Get("dl") == "1" {
				w.Header().Set("Content-Disposition", `attachment; filename="`+strings.ReplaceAll(name, `"`, "")+`"`)
			}
			http.ServeFile(w, r, p)
		default:
			http.NotFound(w, r)
		}
	}))
	return mux
}

var mimeTable = map[string]string{".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp", ".gif": "image/gif",
	".mp4": "video/mp4", ".webm": "video/webm", ".mov": "video/quicktime", ".txt": "text/plain; charset=utf-8",
	".srt": "text/plain; charset=utf-8", ".vtt": "text/vtt; charset=utf-8", ".json": "application/json"}

// mimeFor answers only with types a browser shows as data (never html or script), whatever the file is called.
func mimeFor(name string) string {
	if m := mimeTable[strings.ToLower(filepath.Ext(name))]; m != "" {
		return m
	}
	return "application/octet-stream"
}

// submitMultipart takes a multipart form: capability, input (JSON) and one file part per file field. The files are
// written to disk as they arrive. owner is "studio" or "key:<id>".
func (s *Server) submitMultipart(w http.ResponseWriter, r *http.Request, owner string) (*Job, *apiErr) {
	r.Body = http.MaxBytesReader(w, r.Body, 256<<20)
	mr, err := r.MultipartReader()
	if err != nil {
		return nil, &apiErr{400, "send the form as multipart", "invalid_request"}
	}
	id := randHex(8)
	dir := filepath.Join(s.jobs.jobDir(id), "in")
	bad := func(code int, typ, msg string) (*Job, *apiErr) {
		os.RemoveAll(s.jobs.jobDir(id))
		return nil, &apiErr{code, msg, typ}
	}
	var capID string
	input := map[string]any{}
	files := map[string]string{} // field -> stored name
	var cp modules.Capability
	for {
		part, err := mr.NextPart()
		if err == io.EOF {
			break
		}
		if err != nil {
			return bad(400, "invalid_request", "bad form")
		}
		switch field := part.FormName(); {
		case field == "capability":
			b, _ := io.ReadAll(io.LimitReader(part, 100))
			capID = strings.TrimSpace(string(b))
			var ok bool
			if cp, ok = s.jobs.cat.Capability(capID); !ok {
				return bad(400, "invalid_request", "unknown capability")
			}
		case field == "input":
			b, _ := io.ReadAll(io.LimitReader(part, 64<<10))
			if json.Unmarshal(b, &input) != nil {
				return bad(400, "invalid_json", "input must be JSON")
			}
		case part.FileName() != "":
			if capID == "" {
				return bad(400, "invalid_request", "send capability first")
			}
			name := safeName(part.FileName())
			if name == "" || !isFileField(cp, field) {
				return bad(400, "invalid_request", "unexpected file")
			}
			os.MkdirAll(dir, 0750)
			stored := field + "-" + name
			n, err := saveLimited(filepath.Join(dir, stored), part, cp.MaxInput)
			if err != nil {
				return bad(413, "request_too_large", "the file is too large for this capability")
			}
			if n == 0 {
				return bad(400, "invalid_request", "empty file")
			}
			files[field] = stored
		}
	}
	if capID == "" {
		return bad(400, "invalid_request", "capability is required")
	}
	for f, name := range files {
		input[f] = name
	}
	return s.admitJob(owner, cp, id, input)
}

func isFileField(cp modules.Capability, name string) bool {
	for _, f := range cp.Inputs {
		if f.Name == name && (f.Type == "image" || f.Type == "video" || f.Type == "audio") {
			return true
		}
	}
	return false
}

func saveLimited(path string, r io.Reader, max int64) (int64, error) {
	f, err := os.Create(path)
	if err != nil {
		return 0, err
	}
	defer f.Close()
	if max <= 0 {
		max = 20 << 20
	}
	n, err := io.Copy(f, io.LimitReader(r, max+1))
	if err == nil && n > max {
		err = io.ErrShortBuffer
	}
	if err != nil {
		os.Remove(path)
	}
	return n, err
}