main.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"unicode/utf8"
)

var (
	dmKeywordRe = regexp.MustCompile(`(?m)^\s*(mob|obj|turf|var|world|#define|proc|verb)\b`)

	// Command-line flags
	inputFile    = flag.String("input", "", "Input file")
	inputFolder  = flag.String("input-folder", "", "Input folder")
	outputFolder = flag.String("output", "", "Output file")
	maxBlockSize = flag.Int("max-block-size", 1024*1024, "Max block size for DMI extraction")
	minStringLen = flag.Int("min-string-len", 4, "Minimum string length for text extraction")

	// Choices for extraction
	extractDMIs  = flag.Bool("extract-dmi", true, "Extract DMIs")
	extractDM    = flag.Bool("extract-dm", true, "Extract DM")
	extractText  = flag.Bool("extract-text", true, "Extract text")
	extractPNGs  = flag.Bool("extract-png", true, "Extract PNGs")
	extractAudio = flag.Bool("extract-audio", true, "Extract audio files (MIDI/WAV)")
)

func main() {
	flag.Parse()
	if *inputFile == "" && *inputFolder == "" {
		log.Fatal("Provide an input file using the -input flag.")
	}
	if *inputFile != "" && *inputFolder != "" {
		log.Fatal("Provide either an input file or an input folder, not both.")
	}
	if *outputFolder == "" {
		log.Fatal("Please provide an output folder using the -output flag.")
	}

	// Create the output directory if it doesn't exist
	if err := os.MkdirAll(*outputFolder, 0755); err != nil {
		log.Fatal("failed to create output folder: ", err)
	}

	if inputFile != nil && *inputFile != "" {
		err := processFile(*inputFile, *outputFolder, *extractDMIs, *extractPNGs, *extractText, *extractDM, *extractAudio)
		if err != nil {
			log.Fatal("Error processing file: ", err)
		}
	} else {
		err := processFolder(*inputFolder, *outputFolder, *extractDMIs, *extractPNGs, *extractText, *extractDM, *extractAudio)
		if err != nil {
			log.Fatal("Error processing folder: ", err)
		}
	}
}

// Processes a folder, extracting chosen elements from each file
func processFolder(folderPath string, outputFolder string, extractDMIs, extractPNGs, extractText, extractDM, extractAudio bool) error {
	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("failed to access path %q: %w", path, err)
		}
		if info.IsDir() {
			return nil // Skip directories
		}
		log.Printf("Processing file: %s", path)
		return processFile(path, outputFolder, extractDMIs, extractPNGs, extractText, extractDM, extractAudio)
	})
	if err != nil {
		return fmt.Errorf("error walking the path %q: %w", folderPath, err)
	}
	return nil
}

// Processes a single file, extracting chosen elements
func processFile(filePath string, outputPath string, extractDMIs, extractPNGs, extractText, extractDM, extractAudio bool) error {
	baseFilename := filepath.Base(filePath)

	// Read the input file
	file, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("failed to open input file: %w", err)
	}
	defer file.Close()

	// Extract old-style DMIs, v3
	if extractDMIs {
		err = extractOldDMIs(file, outputPath, *maxBlockSize)
		if err != nil {
			return fmt.Errorf("failed to extract DMIs: %w", err)
		}
	}

	// Extract PNGs
	if extractPNGs {
		extractAllPNGs(file, outputPath)
	}

	// Extract DM files
	if extractDM {
		err = extractDMFiles(file, outputPath)
		if err != nil {
			return fmt.Errorf("failed to extract DM files: %w", err)
		}
	}

	// Extract MIDI files
	if extractAudio {
		extractAllMIDIs(file, outputPath)
		extractAllWavs(file, outputPath)
	}

	// Extract text
	if extractText {
		err = extractTextWithWhitespace(file, fmt.Sprintf("%s/%s.txt", outputPath, trimAllExtensions(baseFilename)), *minStringLen)
		if err != nil {
			return fmt.Errorf("failed to extract text: %w", err)
		}
	}
	return nil
}

// Helper to extract a single PNG from a given offset
func extractSinglePNG(file *os.File, startOffset int64, outPath string) bool {
	const pngSig = "\x89PNG\r\n\x1a\n"
	out, err := os.Create(outPath)
	if err != nil {
		log.Printf("Failed to create output file: %v", err)
		return false
	}
	defer out.Close()

	// Write the PNG signature
	_, err = out.Write([]byte(pngSig))
	if err != nil {
		log.Printf("Failed to write PNG signature: %v", err)
		return false
	}

	offset := startOffset + 8
	for {
		// Read chunk length (4 bytes, big-endian)
		var lengthBuf [4]byte
		if _, err := file.ReadAt(lengthBuf[:], offset); err != nil {
			log.Printf("Failed to read chunk length at %d: %v", offset, err)
			return false
		}
		offset += 4
		length := int64(uint32(lengthBuf[0])<<24 | uint32(lengthBuf[1])<<16 | uint32(lengthBuf[2])<<8 | uint32(lengthBuf[3]))

		// Read chunk type (4 bytes)
		var typeBuf [4]byte
		if _, err := file.ReadAt(typeBuf[:], offset); err != nil {
			log.Printf("Failed to read chunk type at %d: %v", offset, err)
			return false
		}
		offset += 4

		// Write length and type
		out.Write(lengthBuf[:])
		out.Write(typeBuf[:])

		// Read chunk data
		chunkData := make([]byte, length+4) // data + CRC
		if _, err := file.ReadAt(chunkData, offset); err != nil {
			log.Printf("Failed to read chunk data at %d: %v", offset, err)
			return false
		}
		out.Write(chunkData)
		offset += length + 4

		// Check for IEND
		if string(typeBuf[:]) == "IEND" {
			break
		}
	}
	return true
}

// Extracts text from a file, looking for printable ASCII characters and whitespace
func extractTextWithWhitespace(file *os.File, outputPath string, minStringLen int) error {
	const chunkSize = 4 * 1024 * 1024
	buffer := make([]byte, chunkSize)
	out, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer out.Close()

	var strBuilder strings.Builder
	var offset int64 = 0

	isTextChar := func(b byte) bool {
		// Accept printable ASCII, tab, newline, carriage return
		return (b >= 32 && b <= 126) || b == '\n' || b == '\r' || b == '\t'
	}

	for {
		n, err := file.ReadAt(buffer, offset)
		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 {
			break
		}
		i := 0
		for i < n {
			// Skip padding (runs of 0x00)
			if buffer[i] == 0x00 {
				if strBuilder.Len() >= minStringLen {
					out.WriteString(strBuilder.String())
					out.WriteString("\n")
				}
				strBuilder.Reset()
				for i < n && buffer[i] == 0x00 {
					i++
				}
				continue
			}
			// Accept text chars and whitespace
			if isTextChar(buffer[i]) {
				strBuilder.WriteByte(buffer[i])
			} else {
				// Non-text, flush if long enough
				if strBuilder.Len() >= minStringLen {
					out.WriteString(strBuilder.String())
					out.WriteString("\n")
				}
				strBuilder.Reset()
			}
			i++
		}
		offset += int64(n)
	}
	// Write any trailing string
	if strBuilder.Len() >= minStringLen {
		out.WriteString(strBuilder.String())
		out.WriteString("\n")
	}
	log.Printf("Extracted txt file: %s", outputPath)
	return nil
}

// Extracts old-style DMIs from a file, looking for the "\x04DMI" signature(Pre-BYOND 4.0)
func extractOldDMIs(file *os.File, outputDir string, maxBlockSize int) error {
	const chunkSize = 4 * 1024 * 1024
	filenamePattern := regexp.MustCompile(`([ -~]{1,64}\.dmi)\x00\x04DMI`)
	rawPattern := regexp.MustCompile(`\x04DMI`)

	found := make(map[string]int)
	processed := make(map[int64]bool) // Track offsets we've already processed
	buffer := make([]byte, chunkSize)
	var offset int64 = 0
	var rawDmiCount int = 0

	for {
		n, err := file.ReadAt(buffer, offset)
		if err != nil && err != io.EOF {
			return fmt.Errorf("read failed: %w", err)
		}
		if n == 0 {
			break
		}

		chunk := buffer[:n]

		// First pass: Extract DMIs with filename
		locs := filenamePattern.FindAllSubmatchIndex(chunk, -1)
		for _, loc := range locs {
			if len(loc) < 4 {
				continue
			}

			// Extract filename
			startIdx := loc[2]
			endIdx := loc[3]
			filename := string(chunk[startIdx:endIdx])
			filename = filepath.Base(filename)

			// Create a lowercase key for case-insensitive tracking
			lowerKey := strings.ToLower(filename)

			// Check if we've seen this filename (case-insensitive)
			if found[lowerKey] > 0 {
				// Add a number to make it unique
				filename = fmt.Sprintf("%s_%d.dmi", filename[:len(filename)-4], found[lowerKey]+1)
			}
			found[lowerKey] += 1

			// Offset in file where \x04DMI starts
			dmiOffset := offset + int64(loc[1]-4)
			processed[dmiOffset] = true

			// Extract DMI data
			extractDMIData(file, dmiOffset, maxBlockSize, outputDir, filename)
		}

		// Second pass: Extract DMIs without filename
		locs = rawPattern.FindAllIndex(chunk, -1)
		for _, loc := range locs {
			dmiOffset := offset + int64(loc[0])

			// Skip if we already processed this offset
			if processed[dmiOffset] {
				continue
			}
			processed[dmiOffset] = true

			rawDmiCount++
			filename := fmt.Sprintf("unnamed_dmi_%05d.dmi", rawDmiCount)

			// Extract DMI data
			extractDMIData(file, dmiOffset, maxBlockSize, outputDir, filename)
		}

		offset += int64(n)
		if err == io.EOF {
			break
		}
	}

	log.Printf("Total DMIs extracted: %d named, %d unnamed", len(found), rawDmiCount)
	return nil
}

// Helper function to extract DMI data from a given offset
func extractDMIData(file *os.File, dmiOffset int64, maxBlockSize int, outputDir string, filename string) {
	// First, verify the DMI signature
	header := make([]byte, 4)
	_, err := file.ReadAt(header, dmiOffset)
	if err != nil && err != io.EOF {
		log.Printf("Failed to read DMI header at offset %d: %v", dmiOffset, err)
		return
	}
	if string(header) != "\x04DMI" {
		log.Printf("Invalid DMI signature at offset %d", dmiOffset)
		return
	}

	// Read a full chunk (we'll analyze it, not auto-trim)
	buffer := make([]byte, maxBlockSize)
	nRead, err := file.ReadAt(buffer, dmiOffset)
	if err != nil && err != io.EOF {
		log.Printf("Failed to read DMI data at offset %d: %v", dmiOffset, err)
		return
	}

	// FIXED APPROACH: Check for the next DMI or PNG header as a boundary
	nextHeaderPos := nRead
	for i := 100; i < nRead-4; i++ { // Skip first 100 bytes to avoid finding our own header
		// Look for next DMI signature (\x04DMI) or PNG signature (\x89PNG...)
		if (buffer[i] == 0x04 && buffer[i+1] == 'D' && buffer[i+2] == 'M' && buffer[i+3] == 'I') ||
			(buffer[i] == 0x89 && buffer[i+1] == 'P' && buffer[i+2] == 'N' && buffer[i+3] == 'G') {
			// Found another file start
			nextHeaderPos = i
			break
		}
	}

	// Write the file up to the next header or the full buffer if no next header found
	outputPath := filepath.Join(outputDir, filename)
	err = os.WriteFile(outputPath, buffer[:nextHeaderPos], 0644)
	if err != nil {
		log.Printf("Failed to write %s: %v", filename, err)
	} else {
		log.Printf("Extracted DMI file: %s", outputPath)
	}
}

// Checks if a PNG file contains the DMI signature in the first 64KB
func isDMIPNG(file *os.File, offset int64) bool {
	const previewSize = 64 * 1024
	buf := make([]byte, previewSize)
	n, _ := file.ReadAt(buf, offset)
	return n > 0 && strings.Contains(string(buf[:n]), "DMI")
}

// Extracts all PNGs from a file, looking for the "\x89PNG" signature
// and checking for the DMI signature in the first 64KB of the PNG data
func extractAllPNGs(file *os.File, outputDir string) {
	const pngSig = "\x89PNG\r\n\x1a\n"
	const scanChunk = 4 * 1024 * 1024

	buffer := make([]byte, scanChunk)
	var offset int64 = 0
	pngCount := 0

	for {
		n, err := file.ReadAt(buffer, offset)
		if err != nil && err != io.EOF {
			log.Printf("Error reading file: %v", err)
			break
		}
		if n == 0 {
			break
		}
		chunk := buffer[:n]
		for i := 0; i < len(chunk)-8; i++ {
			if string(chunk[i:i+8]) == pngSig {
				pngOffset := offset + int64(i)
				ext := ".png"
				if isDMIPNG(file, pngOffset) {
					ext = ".dmi"
				}
				outPath := filepath.Join(outputDir, fmt.Sprintf("unnamed_png_%05d%s", pngCount, ext))
				if extractSinglePNG(file, pngOffset, outPath) {
					log.Printf("Extracted PNG file: %s", outPath)
					pngCount++
				}
				i += 8 // Skip ahead to avoid double-counting
			}
		}
		offset += int64(n)
	}
	log.Printf("Total PNGs extracted: %d", pngCount)
}

// Extracts DM files from a given file, looking for the DM keyword
func extractDMFiles(f *os.File, outDir string) error {
	reader := bufio.NewReader(f)
	const maxBlockBytes = 200 * 1024
	const minKeywordHits = 3
	const maxNonPrintable = 4
	const preBufferLines = 5

	var block []string
	var prebuffer []string
	var blockSize int
	var keywordHits int
	var nonPrintableStreak int
	var fileCount int
	var insideBlock bool

	flush := func() error {
		if keywordHits >= minKeywordHits && len(block) > 0 {
			outPath := filepath.Join(outDir, fmt.Sprintf("unnamed_dm_%05d.dm", fileCount+1))
			fp, err := os.Create(outPath)
			if err != nil {
				return err
			}
			for _, line := range block {
				_, _ = fp.WriteString(line)
			}
			_ = fp.Close()
			log.Printf("Extracted DM file: %s", outPath)
			fileCount++
		}
		block = nil
		blockSize = 0
		keywordHits = 0
		nonPrintableStreak = 0
		insideBlock = false
		prebuffer = nil
		return nil
	}

	for {
		line, err := reader.ReadString('\n')

		if len(line) > 0 {
			if isLinePrintable(line) {
				nonPrintableStreak = 0

				// Track prebuffer before block begins
				if !insideBlock {
					prebuffer = append(prebuffer, line)
					if len(prebuffer) > preBufferLines {
						prebuffer = prebuffer[1:] // keep only the last few
					}
				}

				if dmKeywordRe.MatchString(line) {
					if !insideBlock {
						block = append(block, prebuffer...)
						insideBlock = true
					}
					keywordHits++
				}

				if insideBlock {
					block = append(block, line)
					blockSize += len(line)
				}
			} else {
				nonPrintableStreak++
			}
		}

		if nonPrintableStreak >= maxNonPrintable || blockSize >= maxBlockBytes {
			if err := flush(); err != nil {
				return err
			}
		}

		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("read error: %w", err)
		}
	}

	// Final flush
	if len(block) > 0 {
		if err := flush(); err != nil {
			return err
		}
	}

	log.Printf("Total DMs extracted: %d", fileCount)

	return nil
}

func isLinePrintable(s string) bool {
	for len(s) > 0 {
		r, size := utf8.DecodeRuneInString(s)
		if r == utf8.RuneError || (!runeIsPrintable(r) && r != '\t' && r != '\n' && r != '\r') {
			return false
		}
		s = s[size:]
	}
	return true
}

func runeIsPrintable(r rune) bool {
	return (r >= 0x20 && r <= 0x7E) || (r >= 0xA1 && r != 0x7F)
}

func trimAllExtensions(path string) string {
	base := filepath.Base(path)
	for ext := filepath.Ext(base); ext != ""; ext = filepath.Ext(base) {
		base = base[:len(base)-len(ext)]
	}
	return base
}

// Extracts all MIDI files from a file, looking for the "MThd" signature and a possible filename before it
func extractAllMIDIs(file *os.File, outputDir string) {
	const midiSig = "MThd"
	const scanChunk = 4 * 1024 * 1024
	filenamePattern := regexp.MustCompile(`([ -~]{1,64}\.mid)\x00MThd`)

	buffer := make([]byte, scanChunk)
	var offset int64 = 0
	midiCount := 0
	usedNames := make(map[string]int)

	for {
		n, err := file.ReadAt(buffer, offset)
		if err != nil && err != io.EOF {
			log.Printf("Error reading file: %v", err)
			break
		}
		if n == 0 {
			break
		}
		chunk := buffer[:n]

		// First pass: look for filename.mid\x00MThd
		locs := filenamePattern.FindAllSubmatchIndex(chunk, -1)
		for _, loc := range locs {
			if len(loc) < 4 {
				continue
			}
			startIdx := loc[2]
			endIdx := loc[3]
			filename := string(chunk[startIdx:endIdx])
			filename = filepath.Base(filename)
			lowerKey := strings.ToLower(filename)
			if usedNames[lowerKey] > 0 {
				filename = fmt.Sprintf("%s_%d.mid", filename[:len(filename)-4], usedNames[lowerKey]+1)
			}
			usedNames[lowerKey]++
			midiOffset := offset + int64(loc[1]-4)
			outPath := filepath.Join(outputDir, filename)
			if extractSingleMIDI(file, midiOffset, outPath) {
				log.Printf("Extracted MIDI file: %s", outPath)
				midiCount++
			}
		}

		// Second pass: generic unnamed MIDIs
		for i := 0; i < len(chunk)-4; i++ {
			if string(chunk[i:i+4]) == midiSig {
				midiOffset := offset + int64(i)
				// Skip if already processed in the filename pass
				alreadyProcessed := false
				for _, loc := range locs {
					if int64(loc[1]-4) == midiOffset {
						alreadyProcessed = true
						break
					}
				}
				if alreadyProcessed {
					continue
				}
				outPath := filepath.Join(outputDir, fmt.Sprintf("unnamed_midi_%05d.mid", midiCount+1))
				if extractSingleMIDI(file, midiOffset, outPath) {
					log.Printf("Extracted MIDI file: %s", outPath)
					midiCount++
				}
				i += 4
			}
		}
		offset += int64(n)
	}
	log.Printf("Total MIDIs extracted: %d", midiCount)
}

// Helper to extract a single MIDI file from a given offset
func extractSingleMIDI(file *os.File, startOffset int64, outPath string) bool {
	// Read MIDI header (14 bytes: "MThd" + 6 bytes header + 4 bytes chunk type)
	header := make([]byte, 14)
	_, err := file.ReadAt(header, startOffset)
	if err != nil {
		log.Printf("Failed to read MIDI header at offset %d: %v", startOffset, err)
		return false
	}
	if string(header[:4]) != "MThd" {
		return false
	}
	// MIDI header always 6 bytes after "MThd"
	// Now, scan for the next "MThd" or EOF as the end of this MIDI
	const maxMIDI = 2 * 1024 * 1024 // 2MB max MIDI size
	buffer := make([]byte, maxMIDI)
	n, err := file.ReadAt(buffer, startOffset)
	if err != nil && err != io.EOF {
		log.Printf("Failed to read MIDI data at offset %d: %v", startOffset, err)
		return false
	}
	end := n
	for i := 4; i < n-4; i++ {
		if string(buffer[i:i+4]) == "MThd" {
			end = i
			break
		}
	}
	err = os.WriteFile(outPath, buffer[:end], 0644)
	if err != nil {
		log.Printf("Failed to write MIDI file: %v", err)
		return false
	}
	return true
}

// Extracts all WAV files from a file, looking for the "RIFF" signature and a possible filename before it
func extractAllWavs(file *os.File, outputDir string) {
	const wavSig = "RIFF"
	const scanChunk = 4 * 1024 * 1024
	filenamePattern := regexp.MustCompile(`([ -~]{1,64}\.wav)\x00RIFF`)

	buffer := make([]byte, scanChunk)
	var offset int64 = 0
	wavCount := 0
	usedNames := make(map[string]int)

	for {
		n, err := file.ReadAt(buffer, offset)
		if err != nil && err != io.EOF {
			log.Printf("Error reading file: %v", err)
			break
		}
		if n == 0 {
			break
		}
		chunk := buffer[:n]

		// First pass: look for filename.wav\x00RIFF
		locs := filenamePattern.FindAllSubmatchIndex(chunk, -1)
		for _, loc := range locs {
			if len(loc) < 4 {
				continue
			}
			startIdx := loc[2]
			endIdx := loc[3]
			filename := string(chunk[startIdx:endIdx])
			filename = filepath.Base(filename)
			lowerKey := strings.ToLower(filename)
			if usedNames[lowerKey] > 0 {
				filename = fmt.Sprintf("%s_%d.wav", filename[:len(filename)-4], usedNames[lowerKey]+1)
			}
			usedNames[lowerKey]++
			wavOffset := offset + int64(loc[1]-4)
			outPath := filepath.Join(outputDir, filename)
			if extractSingleWAV(file, wavOffset, outPath) {
				log.Printf("Extracted WAV file: %s", outPath)
				wavCount++
			}
		}

		// Second pass: generic unnamed WAVs
		for i := 0; i < len(chunk)-4; i++ {
			if string(chunk[i:i+4]) == wavSig {
				wavOffset := offset + int64(i)
				// Skip if already processed in the filename pass
				alreadyProcessed := false
				for _, loc := range locs {
					if int64(loc[1]-4) == wavOffset {
						alreadyProcessed = true
						break
					}
				}
				if alreadyProcessed {
					continue
				}
				outPath := filepath.Join(outputDir, fmt.Sprintf("unnamed_wav_%05d.wav", wavCount+1))
				if extractSingleWAV(file, wavOffset, outPath) {
					log.Printf("Extracted WAV file: %s", outPath)
					wavCount++
				}
				i += 4
			}
		}
		offset += int64(n)
	}
	log.Printf("Total WAVs extracted: %d", wavCount)
}

// Helper to extract a single WAV file from a given offset
func extractSingleWAV(file *os.File, startOffset int64, outPath string) bool {
	// Read WAV header (44 bytes: "RIFF" + 4 bytes header + 4 bytes chunk type)
	header := make([]byte, 44)
	_, err := file.ReadAt(header, startOffset)
	if err != nil {
		log.Printf("Failed to read WAV header at offset %d: %v", startOffset, err)
		return false
	}
	if string(header[:4]) != "RIFF" {
		return false
	}
	// WAV header always 36 bytes after "RIFF"
	// Now, scan for the next "RIFF" or EOF as the end of this WAV
	const maxWAV = 2 * 1024 * 1024 // 2MB max WAV size
	buffer := make([]byte, maxWAV)
	n, err := file.ReadAt(buffer, startOffset)
	if err != nil && err != io.EOF {
		log.Printf("Failed to read WAV data at offset %d: %v", startOffset, err)
		return false
	}
	end := n
	for i := 4; i < n-4; i++ {
		if string(buffer[i:i+4]) == "RIFF" {
			end = i
			break
		}
	}
	err = os.WriteFile(outPath, buffer[:end], 0644)
	if err != nil {
		log.Printf("Failed to write WAV file: %v", err)
		return false
	}
	return true
}