package sanitize

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"path"
	"strings"

	"git.mchus.pro/mchus/logpile/internal/privacy"
)

const (
	maxInputBytes  = 800 << 20 // whole file, read into memory and rebuilt
	maxMemberBytes = 1 << 30   // single decompressed member
)

// rewriteBytes de-identifies data according to its format (by filename
// extension), recursing into nested archives. It returns the rebuilt bytes plus
// the replacements made and any binary members that held customer data but
// could not be edited safely.
func rewriteBytes(name string, data []byte) (out []byte, changes []change, skippedBinary []string, err error) {
	switch strings.ToLower(path.Ext(name)) {
	case ".tar", ".sds":
		return rewriteTar(name, data)
	case ".gz", ".tgz":
		return rewriteGz(name, data)
	case ".zip":
		return rewriteZip(name, data)
	case ".ahs":
		// Proprietary HPE container - not editable in place. Report if it holds
		// customer data.
		if scanHasFindings(data) {
			return data, nil, []string{name + " (HPE AHS container - sanitize manually)"}, nil
		}
		return data, nil, nil, nil
	case ".txt", ".log":
		if privacy.IsAllowlistedFile(name) {
			return data, nil, nil, nil
		}
		nb, ch := redactText(data, isCertName(name))
		return nb, prefixPath(ch, path.Base(name)), nil, nil
	default:
		return rewriteMember(name, data, false)
	}
}

// rewriteMember handles one file inside an archive: recurse if it is itself an
// archive, redact if it is text, otherwise leave it (and flag it when a binary
// member carries customer data). Returned changes/skips are already qualified
// with name, so callers append them verbatim.
func rewriteMember(name string, data []byte, isDirOrSpecial bool) ([]byte, []change, []string, error) {
	if isDirOrSpecial || len(data) == 0 || privacy.IsAllowlistedFile(name) {
		return data, nil, nil, nil
	}
	if isNestedArchive(name) || looksLikeTar(data) {
		nb, ch, skip, err := rewriteBytes(name, data)
		if err == nil {
			return nb, prefixPath(ch, name), prefixNames(skip, name), nil
		}
		// A truncated / mis-named "archive" must not fail the whole job: treat it
		// as a plain file if it is text, otherwise copy it and flag it.
		if looksLikeText(data) {
			nb, ch := redactText(data, isCertName(name))
			return nb, prefixPath(ch, name), nil, nil
		}
		if len(data) <= 8<<20 && binaryHasLeak(data) {
			return data, nil, []string{name + " (unreadable as an archive; holds customer data - sanitize manually)"}, nil
		}
		return data, nil, nil, nil
	}
	if looksLikeText(data) {
		nb, ch := redactText(data, isCertName(name))
		return nb, prefixPath(ch, name), nil, nil
	}
	// Binary member: never edit it (checksums / structure), but tell the
	// operator if a printable run inside it holds customer data.
	if len(data) <= 8<<20 && binaryHasLeak(data) {
		return data, nil, []string{name + " (binary - sanitize manually)"}, nil
	}
	return data, nil, nil, nil
}

func rewriteTar(name string, data []byte) ([]byte, []change, []string, error) {
	tr := tar.NewReader(bytes.NewReader(data))
	var buf bytes.Buffer
	tw := tar.NewWriter(&buf)
	var changes []change
	var skipped []string

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, nil, nil, fmt.Errorf("%s: tar read: %w", name, err)
		}
		body, err := io.ReadAll(io.LimitReader(tr, maxMemberBytes+1))
		if err != nil {
			return nil, nil, nil, fmt.Errorf("%s: read %s: %w", name, hdr.Name, err)
		}
		special := !hdr.FileInfo().Mode().IsRegular()
		newBody, ch, skip, err := rewriteMember(hdr.Name, body, special || int64(len(body)) > maxMemberBytes)
		if err != nil {
			return nil, nil, nil, err
		}
		h := *hdr // reuse every header field verbatim (name, mode, uid/gid, mtime/atime/ctime, pax, format)
		h.Size = int64(len(newBody))
		if err := tw.WriteHeader(&h); err != nil {
			return nil, nil, nil, fmt.Errorf("%s: write header %s: %w", name, hdr.Name, err)
		}
		if _, err := tw.Write(newBody); err != nil {
			return nil, nil, nil, err
		}
		changes = append(changes, ch...)
		skipped = append(skipped, skip...)
	}
	if len(changes) == 0 {
		return data, nil, skipped, nil // nothing redacted -> byte-identical
	}
	if err := tw.Close(); err != nil {
		return nil, nil, nil, err
	}
	return buf.Bytes(), changes, skipped, nil
}

func rewriteGz(name string, data []byte) ([]byte, []change, []string, error) {
	gzr, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, nil, nil, fmt.Errorf("%s: gzip: %w", name, err)
	}
	hdr := gzr.Header
	decompressed, err := io.ReadAll(io.LimitReader(gzr, maxMemberBytes+1))
	gzr.Close()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("%s: gunzip: %w", name, err)
	}
	if int64(len(decompressed)) > maxMemberBytes {
		return data, nil, []string{name + " (too large to sanitize)"}, nil
	}

	innerName := strings.TrimSuffix(hdr.Name, ".gz")
	if innerName == "" {
		innerName = strings.TrimSuffix(path.Base(name), ".gz")
	}

	var newInner []byte
	var ch []change
	var skip []string
	switch {
	case looksLikeTar(decompressed):
		newInner, ch, skip, err = rewriteTar(innerName, decompressed)
		if err != nil {
			return data, nil, []string{name + " (unreadable inner tar - copied as-is)"}, nil
		}
	case privacy.IsAllowlistedFile(innerName):
		newInner = decompressed
	case looksLikeText(decompressed):
		newInner, ch = redactText(decompressed, isCertName(innerName))
		ch = prefixPath(ch, innerName)
	case len(decompressed) <= 8<<20 && binaryHasLeak(decompressed):
		newInner, skip = decompressed, []string{innerName + " (binary - sanitize manually)"}
	default:
		newInner = decompressed
	}
	if len(ch) == 0 {
		return data, nil, skip, nil // nothing redacted -> byte-identical
	}

	var buf bytes.Buffer
	gzw, _ := gzip.NewWriterLevel(&buf, gzip.BestCompression)
	gzw.Name = hdr.Name
	gzw.Comment = hdr.Comment
	gzw.ModTime = hdr.ModTime
	gzw.OS = hdr.OS
	gzw.Extra = hdr.Extra
	if _, err := gzw.Write(newInner); err != nil {
		return nil, nil, nil, err
	}
	if err := gzw.Close(); err != nil {
		return nil, nil, nil, err
	}
	return buf.Bytes(), ch, skip, nil
}

func rewriteZip(name string, data []byte) ([]byte, []change, []string, error) {
	zr, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
	if err != nil {
		return nil, nil, nil, fmt.Errorf("%s: zip: %w", name, err)
	}
	var buf bytes.Buffer
	zw := zip.NewWriter(&buf)
	if zr.Comment != "" {
		_ = zw.SetComment(zr.Comment)
	}
	var changes []change
	var skipped []string

	for _, f := range zr.File {
		if f.FileInfo().IsDir() {
			if err := zw.Copy(f); err != nil {
				return nil, nil, nil, err
			}
			continue
		}
		rc, err := f.Open()
		if err == nil {
			var body []byte
			body, err = io.ReadAll(io.LimitReader(rc, maxMemberBytes+1))
			rc.Close()
			if err == nil {
				newBody, ch, skip, mErr := rewriteMember(f.Name, body, int64(len(body)) > maxMemberBytes)
				if mErr == nil {
					if bytes.Equal(newBody, body) {
						if err := zw.Copy(f); err != nil {
							return nil, nil, nil, err
						}
					} else {
						fh := f.FileHeader
						fh.CRC32, fh.CompressedSize, fh.CompressedSize64 = 0, 0, 0
						fh.UncompressedSize, fh.UncompressedSize64 = 0, 0
						w, cErr := zw.CreateHeader(&fh)
						if cErr != nil {
							return nil, nil, nil, cErr
						}
						if _, wErr := w.Write(newBody); wErr != nil {
							return nil, nil, nil, wErr
						}
					}
					changes = append(changes, ch...)
					skipped = append(skipped, skip...)
					continue
				}
			}
		}
		// Unreadable / unprocessable entry: copy it raw, flag it.
		if err := zw.Copy(f); err != nil {
			return nil, nil, nil, fmt.Errorf("%s: copy %s: %w", name, f.Name, err)
		}
		skipped = append(skipped, f.Name+" (unreadable zip entry - copied as-is)")
	}
	if len(changes) == 0 {
		return data, nil, skipped, nil // nothing redacted -> byte-identical
	}
	if err := zw.Close(); err != nil {
		return nil, nil, nil, err
	}
	return buf.Bytes(), changes, skipped, nil
}

func isNestedArchive(name string) bool {
	switch strings.ToLower(path.Ext(name)) {
	case ".gz", ".tgz", ".tar", ".zip", ".sds":
		return true
	}
	return false
}

func isCertName(name string) bool {
	switch strings.ToLower(path.Ext(name)) {
	case ".pem", ".csr", ".crt", ".cer":
		return true
	}
	return false
}

func looksLikeTar(b []byte) bool {
	if len(b) < 512 {
		return false
	}
	_, err := tar.NewReader(bytes.NewReader(b)).Next()
	return err == nil
}

// looksLikeText mirrors the privacy scanner's heuristic: no NUL byte in the
// first 8 KiB.
func looksLikeText(b []byte) bool {
	if len(b) == 0 {
		return false
	}
	head := b
	if len(head) > 8192 {
		head = head[:8192]
	}
	return bytes.IndexByte(head, 0) < 0
}

func scanHasFindings(b []byte) bool {
	rep := privacy.Scan([]privacy.File{{Path: "member", Content: b}})
	return rep != nil && rep.Summary.Total > 0
}

// binaryHasLeak reports whether any printable ASCII run inside a binary member
// contains a redactable span (privacy.Scan itself skips non-text files, so it
// cannot see strings embedded in FRU.bin / redis-dump.rdb / SDR.dat).
func binaryHasLeak(b []byte) bool {
	start := -1
	check := func(run []byte) bool {
		if len(run) < 6 {
			return false
		}
		for _, sp := range privacy.FindSpans(string(run), false) {
			if redactableCat[sp.Category] {
				return true
			}
		}
		return false
	}
	for i, c := range b {
		printable := c >= 0x20 && c < 0x7f
		if printable && start < 0 {
			start = i
		}
		if !printable && start >= 0 {
			if check(b[start:i]) {
				return true
			}
			start = -1
		}
	}
	if start >= 0 && check(b[start:]) {
		return true
	}
	return false
}

func prefixPath(ch []change, parent string) []change {
	for i := range ch {
		if ch[i].path == "" {
			ch[i].path = parent
		} else {
			ch[i].path = parent + "::" + ch[i].path
		}
	}
	return ch
}

func prefixNames(names []string, parent string) []string {
	out := make([]string, len(names))
	for i, n := range names {
		out[i] = parent + "::" + n
	}
	return out
}