// stems.go implements audio.StemSeparationProvider against a Demucs shim
// (audioutils) reached through llama-swap's /upstream passthrough (ADR-0024):
//
//	POST /upstream/<id>/v1/stems   multipart file[,model,two_stems,format]
//
// The response is a ZIP of the separated stems — one entry per stem, entry
// name = stem name, extension = container. Zip transport keeps a 4-stem WAV
// result (hundreds of MB decoded) off the JSON-of-base64 path entirely.
package llamaswap

import (
	"archive/zip"
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"path"
	"strings"

	"gitea.stevedudenhoeffer.com/steve/majordomo/audio"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// maxStemsResponseBytes caps the /v1/stems zip body: four WAV stems of a
// long song legitimately pass the 64MB JSON cap. Bounded so a buggy
// upstream can't allocate without limit.
const maxStemsResponseBytes = 512 << 20

// maxStemEntryBytes caps ONE decompressed zip entry — the zip-bomb guard.
// A single WAV stem of even a very long song stays far under this.
const maxStemEntryBytes = 256 << 20

// maxStemEntries caps how many stem entries are unpacked: Demucs emits at
// most six, so anything past a small multiple of that is a hostile or
// broken archive, not a result.
const maxStemEntries = 16

// maxStemsTotalBytes caps the AGGREGATE decompressed size across entries —
// the per-entry bound alone would still let a many-entry bomb multiply up.
const maxStemsTotalBytes = 1 << 30

// StemSeparatorModel implements audio.StemSeparationProvider. The id selects
// which upstream llama-swap loads.
func (p *Provider) StemSeparatorModel(id string, opts ...audio.StemSeparatorModelOption) (audio.StemSeparator, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = audio.ApplyStemSeparatorModelOptions(opts)
	return &stemSeparatorModel{p: p, id: id}, nil
}

type stemSeparatorModel struct {
	p  *Provider
	id string
}

// SeparateStems implements audio.StemSeparator.
func (m *stemSeparatorModel) SeparateStems(ctx context.Context, req audio.StemSeparationRequest, opts ...audio.StemSeparationOption) (*audio.StemSeparationResult, error) {
	req = req.Apply(opts...)
	if len(req.Audio) == 0 {
		return nil, fmt.Errorf("%w: stem separation requires audio bytes", llm.ErrUnsupported)
	}
	mode := strings.ToLower(strings.TrimSpace(req.Mode))
	if mode != "" && mode != "two" && mode != "four" {
		return nil, fmt.Errorf("%w: stem mode must be \"two\" or \"four\", got %q", llm.ErrUnsupported, req.Mode)
	}
	format := strings.ToLower(strings.TrimSpace(req.Format))
	if format != "" && format != "mp3" && format != "wav" {
		return nil, fmt.Errorf("%w: stem format must be \"mp3\" or \"wav\", got %q", llm.ErrUnsupported, req.Format)
	}
	upPath, err := upstreamPath(m.id, "/v1/stems")
	if err != nil {
		return nil, err
	}
	// Demucs' two-stem mode is "isolate one source vs the rest"; vocals is
	// the split this surface promises ("two" = vocals + accompaniment).
	twoStems := ""
	if mode == "two" {
		twoStems = "vocals"
	}
	body, contentType, err := buildMultipart("build stems form",
		filePart{field: "file", filename: transcriptionFilename(req.Filename, req.MIME), data: req.Audio},
		[]formField{
			{"model", req.Model, false},
			{"two_stems", twoStems, false},
			{"format", format, false},
		})
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, upPath, m.id, contentType, body, maxStemsResponseBytes)
	if err != nil {
		return nil, err
	}
	if len(raw) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "stems response contained no data"}
	}
	zr, err := zip.NewReader(bytes.NewReader(raw), int64(len(raw)))
	if err != nil {
		// A non-zip 2xx body is a misconfigured upstream (an HTML error page
		// behind a proxy, a JSON soft error) — fail loud, quoting a slice.
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
			Message: fmt.Sprintf("stems response is not a zip (Content-Type %q): %s", respType, truncateForError(raw))}
	}
	res := &audio.StemSeparationResult{}
	var totalBytes int64
	for _, f := range zr.File {
		if f.FileInfo().IsDir() {
			continue
		}
		if len(res.Stems) >= maxStemEntries {
			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
				Message: fmt.Sprintf("stems zip holds more than %d entries", maxStemEntries)}
		}
		data, err := readZipEntry(f)
		if err != nil {
			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
				Message: fmt.Sprintf("stems zip entry %q: %v", f.Name, err)}
		}
		totalBytes += int64(len(data))
		if totalBytes > maxStemsTotalBytes {
			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
				Message: fmt.Sprintf("stems zip decompresses past %d bytes", int64(maxStemsTotalBytes))}
		}
		// Entry name → stem name; extension → MIME. Entries may sit in a
		// per-model directory ("htdemucs/vocals.mp3"), so use the base name.
		base := path.Base(f.Name)
		ext := path.Ext(base)
		res.Stems = append(res.Stems, audio.Stem{
			Name:  strings.TrimSuffix(base, ext),
			Audio: data,
			MIME:  stemMIME(ext),
		})
	}
	if len(res.Stems) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "stems zip contained no stems"}
	}
	return res, nil
}

// readZipEntry decompresses one entry, bounded by maxStemEntryBytes so a
// zip bomb can't allocate without limit.
func readZipEntry(f *zip.File) ([]byte, error) {
	rc, err := f.Open()
	if err != nil {
		return nil, err
	}
	defer rc.Close()
	data, err := io.ReadAll(io.LimitReader(rc, maxStemEntryBytes+1))
	if err != nil {
		return nil, err
	}
	if int64(len(data)) > maxStemEntryBytes {
		return nil, fmt.Errorf("decompressed entry exceeds %d bytes", int64(maxStemEntryBytes))
	}
	return data, nil
}

// stemMIME maps a stem file extension to its MIME type, reusing speechMIME's
// format table ("" and unknown extensions land on the mp3 default — Demucs'
// own default container).
func stemMIME(ext string) string {
	return speechMIME("", strings.TrimPrefix(strings.ToLower(ext), "."))
}