package main

import (
	"bytes"
	"context"
	"fmt"
	"io/fs"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
)

// Container extensions holding uncompressed PCM. Raw headerless .pcm is
// excluded on purpose: sample rate, width and channel count cannot be
// recovered from the file itself.
var pcmExts = map[string]bool{
	".wav":  true,
	".wave": true,
	".bwf":  true,
	".rf64": true,
	".bw64": true,
	".w64":  true,
	".aif":  true,
	".aiff": true,
	".aifc": true,
}

// findPCM returns every regular file under root whose extension names a
// PCM container. Symlinks are not followed.
func findPCM(root string) ([]string, error) {
	var found []string

	err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.Type().IsRegular() {
			return nil
		}
		if pcmExts[strings.ToLower(filepath.Ext(path))] {
			found = append(found, path)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	sort.Strings(found)
	return found, nil
}

type encoder struct {
	threads int
	dryRun  bool
}

// encodeOne compresses src, verifies the result decodes back to
// identical audio, and only then removes src.
func (e *encoder) encodeOne(ctx context.Context, src string) error {
	codecName, err := probeCodec(ctx, src)
	if err != nil {
		return err
	}

	var dst, codec, level string
	switch {
	case strings.HasPrefix(codecName, "pcm_f32"), strings.HasPrefix(codecName, "pcm_f64"):
		dst, codec, level = replaceExt(src, ".wv"), "wavpack", "3"
	case strings.HasPrefix(codecName, "pcm_"):
		dst, codec, level = replaceExt(src, ".flac"), "flac", "12"
	default:
		slog.Warn("skipping, not PCM", "src", src, "codec", codecName)
		return nil
	}

	if _, err := os.Lstat(dst); err == nil {
		return fmt.Errorf("destination already exists: %s", dst)
	}

	if e.dryRun {
		slog.Info("would encode", "src", src, "dst", dst, "codec", codec)
		return nil
	}

	srcInfo, err := os.Stat(src)
	if err != nil {
		return err
	}

	if err := ffmpeg(ctx,
		"-threads", fmt.Sprint(e.threads),
		"-i", src,
		"-map", "0:a:0",
		"-c:a", codec,
		"-compression_level", level,
		dst,
	); err != nil {
		os.Remove(dst)
		return err
	}

	srcSum, err := audioMD5(ctx, src)
	if err != nil {
		os.Remove(dst)
		return err
	}
	dstSum, err := audioMD5(ctx, dst)
	if err != nil {
		os.Remove(dst)
		return err
	}
	if srcSum != dstSum {
		os.Remove(dst)
		return fmt.Errorf("verification failed, original kept: %s != %s", srcSum, dstSum)
	}

	dstInfo, err := os.Stat(dst)
	if err != nil {
		return err
	}

	if err := os.Remove(src); err != nil {
		return err
	}

	slog.Info("encoded",
		"src", src,
		"dst", dst,
		"ratio", fmt.Sprintf("%.3f", float64(dstInfo.Size())/float64(srcInfo.Size())),
	)
	return nil
}

// probeCodec reports the codec of the file's first audio stream, e.g.
// pcm_s24le or pcm_f32le.
func probeCodec(ctx context.Context, path string) (string, error) {
	cmd := exec.CommandContext(ctx, "ffprobe",
		"-v", "error",
		"-select_streams", "a:0",
		"-show_entries", "stream=codec_name",
		"-of", "default=nw=1:nokey=1",
		path,
	)
	cmd.Stdin = nil

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("ffprobe: %w: %s", err, strings.TrimSpace(stderr.String()))
	}

	codec := strings.TrimSpace(stdout.String())
	if codec == "" {
		return "", fmt.Errorf("no audio stream")
	}
	return codec, nil
}

// audioMD5 hashes the decoded audio, forced to 64-bit float so that
// neither 32-bit integer nor 32-bit float samples lose anything on the
// way to the hash. Without an explicit codec the md5 muxer would
// default to pcm_s16le and hide real differences.
func audioMD5(ctx context.Context, path string) (string, error) {
	cmd := exec.CommandContext(ctx, "ffmpeg",
		"-nostdin", "-v", "error",
		"-i", path,
		"-map", "0:a:0",
		"-c:a", "pcm_f64le",
		"-f", "md5", "-",
	)
	cmd.Stdin = nil

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("ffmpeg md5 of %s: %w: %s", path, err, strings.TrimSpace(stderr.String()))
	}

	sum := strings.TrimSpace(stdout.String())
	if !strings.HasPrefix(sum, "MD5=") {
		return "", fmt.Errorf("unexpected md5 output for %s: %q", path, sum)
	}
	return strings.TrimPrefix(sum, "MD5="), nil
}

func ffmpeg(ctx context.Context, args ...string) error {
	full := append([]string{"-nostdin", "-v", "error"}, args...)

	cmd := exec.CommandContext(ctx, "ffmpeg", full...)
	cmd.Stdin = nil

	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ffmpeg: %w: %s", err, strings.TrimSpace(stderr.String()))
	}
	return nil
}

func replaceExt(path, ext string) string {
	return strings.TrimSuffix(path, filepath.Ext(path)) + ext
}