package kb

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

type StagedUADashboard struct {
	Format   string
	Files    map[string][]byte
	Manifest UADashboardManifest
}

type uaArchiveEntry struct {
	Name      string
	Directory bool
	Read      func() ([]byte, error)
}

type uaArchiveCollector struct {
	ctx       context.Context
	stageDir  string
	limits    UADashboardImportLimits
	files     map[string][]byte
	seen      map[string]struct{}
	rootMode  string
	entries   int
	extracted int64
}

func StageUADashboardArchive(ctx context.Context, archivePath, filename string, limits UADashboardImportLimits) (StagedUADashboard, error) {
	limits = normalizedUADashboardLimits(limits)
	info, err := os.Stat(archivePath)
	if err != nil {
		return StagedUADashboard{}, fmt.Errorf("read dashboard archive: %w", err)
	}
	if !info.Mode().IsRegular() {
		return StagedUADashboard{}, fmt.Errorf("%w: upload is not a regular file", ErrUADashboardArchiveUnsafe)
	}
	if info.Size() > limits.HTTPBytes {
		return StagedUADashboard{}, fmt.Errorf("%w: compressed archive exceeds limit", ErrUADashboardArchiveLimit)
	}
	format, err := detectUADashboardArchiveFormat(archivePath, filename)
	if err != nil {
		return StagedUADashboard{}, err
	}
	stageDir, err := os.MkdirTemp("", "meshx-ua-dashboard-stage-*")
	if err != nil {
		return StagedUADashboard{}, fmt.Errorf("create dashboard staging directory: %w", err)
	}
	defer os.RemoveAll(stageDir)
	if err := os.Chmod(stageDir, 0o700); err != nil {
		return StagedUADashboard{}, fmt.Errorf("secure dashboard staging directory: %w", err)
	}

	files := make(map[string][]byte)
	collector := uaArchiveCollector{ctx: ctx, stageDir: stageDir, limits: limits, files: files, seen: make(map[string]struct{})}
	if format == "zip" {
		err = walkUADashboardZip(archivePath, limits, collector.consume)
	} else {
		err = walkUADashboardTarGz(archivePath, limits, collector.consume)
	}
	if err != nil {
		return StagedUADashboard{}, err
	}
	if !hasExactUADashboardFiles(files) {
		return StagedUADashboard{}, fmt.Errorf("%w: exact dashboard file set is required", ErrUADashboardArchiveLayout)
	}
	manifest, err := validateStagedUADashboard(files, limits)
	if err != nil {
		return StagedUADashboard{}, err
	}
	return StagedUADashboard{Format: format, Files: files, Manifest: manifest}, nil
}

func (collector *uaArchiveCollector) consume(entry uaArchiveEntry) error {
	if err := collector.ctx.Err(); err != nil {
		return err
	}
	collector.entries++
	if collector.entries > collector.limits.Entries {
		return fmt.Errorf("%w: archive entry count exceeds limit", ErrUADashboardArchiveLimit)
	}
	clean, err := normalizeUADashboardArchivePath(entry.Name, entry.Directory)
	if err != nil {
		return err
	}
	if _, exists := collector.seen[clean]; exists {
		return fmt.Errorf("%w: duplicate archive path", ErrUADashboardArchiveUnsafe)
	}
	collector.seen[clean] = struct{}{}
	if entry.Directory {
		if clean != ".ua" {
			return fmt.Errorf("%w: unknown archive directory", ErrUADashboardArchiveLayout)
		}
		if collector.rootMode != "" && collector.rootMode != "ua" {
			return fmt.Errorf("%w: mixed dashboard roots", ErrUADashboardArchiveLayout)
		}
		collector.rootMode = "ua"
		return nil
	}
	mode, basename, ok := uaDashboardArchiveFileLocation(clean)
	if !ok {
		return fmt.Errorf("%w: unknown dashboard file", ErrUADashboardArchiveLayout)
	}
	if collector.rootMode != "" && collector.rootMode != mode {
		return fmt.Errorf("%w: mixed dashboard roots", ErrUADashboardArchiveLayout)
	}
	collector.rootMode = mode
	data, err := entry.Read()
	if err != nil {
		return err
	}
	if int64(len(data)) > collector.limits.FileBytes || collector.extracted > collector.limits.ExtractedBytes-int64(len(data)) {
		return fmt.Errorf("%w: extracted dashboard data exceeds limit", ErrUADashboardArchiveLimit)
	}
	collector.extracted += int64(len(data))
	if err := writeUADashboardStagedFile(collector.stageDir, basename, data); err != nil {
		return err
	}
	collector.files[basename] = data
	return nil
}

func writeUADashboardStagedFile(stageDir, basename string, data []byte) error {
	staged, err := os.OpenFile(filepath.Join(stageDir, basename), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return fmt.Errorf("stage dashboard archive file: %w", err)
	}
	_, writeErr := staged.Write(data)
	closeErr := staged.Close()
	if writeErr != nil {
		return fmt.Errorf("stage dashboard archive file: %w", writeErr)
	}
	if closeErr != nil {
		return fmt.Errorf("close dashboard archive file: %w", closeErr)
	}
	return nil
}

func normalizedUADashboardLimits(limits UADashboardImportLimits) UADashboardImportLimits {
	defaults := DefaultUADashboardImportLimits()
	if limits.HTTPBytes <= 0 {
		limits.HTTPBytes = defaults.HTTPBytes
	}
	if limits.ExtractedBytes <= 0 {
		limits.ExtractedBytes = defaults.ExtractedBytes
	}
	if limits.FileBytes <= 0 {
		limits.FileBytes = defaults.FileBytes
	}
	if limits.Entries <= 0 {
		limits.Entries = defaults.Entries
	}
	if limits.JSONDepth <= 0 {
		limits.JSONDepth = defaults.JSONDepth
	}
	return limits
}

func detectUADashboardArchiveFormat(archivePath, filename string) (string, error) {
	file, err := os.Open(archivePath)
	if err != nil {
		return "", err
	}
	defer file.Close()
	header := make([]byte, 4)
	n, err := io.ReadFull(file, header)
	if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
		return "", err
	}
	header = header[:n]
	format := ""
	if len(header) >= 4 && bytes.Equal(header[:4], []byte{'P', 'K', 3, 4}) {
		format = "zip"
	}
	if len(header) >= 2 && header[0] == 0x1f && header[1] == 0x8b {
		format = "tar.gz"
	}
	if format == "" {
		return "", fmt.Errorf("%w: unsupported archive format", ErrUADashboardArchiveUnsupported)
	}
	lower := strings.ToLower(strings.TrimSpace(filename))
	if lower != "" {
		extension := filepath.Ext(lower)
		compatible := extension == "" || (format == "zip" && extension == ".zip") || (format == "tar.gz" && (strings.HasSuffix(lower, ".tar.gz") || extension == ".tgz"))
		if !compatible {
			return "", fmt.Errorf("%w: filename extension does not match archive format", ErrUADashboardArchiveUnsupported)
		}
	}
	return format, nil
}

func walkUADashboardZip(archivePath string, limits UADashboardImportLimits, consume func(uaArchiveEntry) error) error {
	if err := validateUADashboardZipEnd(archivePath); err != nil {
		return err
	}
	reader, err := zip.OpenReader(archivePath)
	if err != nil {
		return fmt.Errorf("%w: invalid ZIP archive", ErrUADashboardArchiveUnsupported)
	}
	defer reader.Close()
	for _, file := range reader.File {
		if file.Flags&1 != 0 {
			return fmt.Errorf("%w: encrypted ZIP entries are not supported", ErrUADashboardArchiveUnsafe)
		}
		mode := file.Mode()
		directory := file.FileInfo().IsDir()
		if !directory && mode&os.ModeType != 0 {
			return fmt.Errorf("%w: special ZIP entry", ErrUADashboardArchiveUnsafe)
		}
		file := file
		err := consume(uaArchiveEntry{Name: file.Name, Directory: directory, Read: func() ([]byte, error) {
			if file.UncompressedSize64 > uint64(limits.FileBytes) {
				return nil, fmt.Errorf("%w: ZIP entry exceeds limit", ErrUADashboardArchiveLimit)
			}
			reader, err := file.Open()
			if err != nil {
				return nil, fmt.Errorf("%w: invalid ZIP entry", ErrUADashboardArchiveUnsafe)
			}
			defer reader.Close()
			data, err := io.ReadAll(io.LimitReader(reader, limits.FileBytes+1))
			if err != nil {
				return nil, fmt.Errorf("%w: invalid ZIP entry", ErrUADashboardArchiveUnsafe)
			}
			if int64(len(data)) > limits.FileBytes || uint64(len(data)) != file.UncompressedSize64 {
				return nil, fmt.Errorf("%w: ZIP entry size mismatch", ErrUADashboardArchiveLimit)
			}
			return data, nil
		}})
		if err != nil {
			return err
		}
	}
	return nil
}

func validateUADashboardZipEnd(archivePath string) error {
	data, err := os.ReadFile(archivePath)
	if err != nil {
		return err
	}
	start := len(data) - (65535 + 22)
	if start < 0 {
		start = 0
	}
	index := bytes.LastIndex(data[start:], []byte{'P', 'K', 5, 6})
	if index < 0 {
		return fmt.Errorf("%w: ZIP end record is missing", ErrUADashboardArchiveUnsupported)
	}
	index += start
	if index+22 > len(data) {
		return fmt.Errorf("%w: invalid ZIP end record", ErrUADashboardArchiveUnsafe)
	}
	commentBytes := int(binary.LittleEndian.Uint16(data[index+20 : index+22]))
	if index+22+commentBytes != len(data) {
		return fmt.Errorf("%w: trailing ZIP payload", ErrUADashboardArchiveUnsafe)
	}
	return nil
}

func walkUADashboardTarGz(archivePath string, limits UADashboardImportLimits, consume func(uaArchiveEntry) error) error {
	file, err := os.Open(archivePath)
	if err != nil {
		return err
	}
	defer file.Close()
	buffered := bufio.NewReader(file)
	gzipReader, err := gzip.NewReader(buffered)
	if err != nil {
		return fmt.Errorf("%w: invalid gzip archive", ErrUADashboardArchiveUnsupported)
	}
	gzipReader.Multistream(false)
	tarReader := tar.NewReader(gzipReader)
	for {
		header, err := tarReader.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			_ = gzipReader.Close()
			return fmt.Errorf("%w: invalid tar archive", ErrUADashboardArchiveUnsafe)
		}
		directory := header.Typeflag == tar.TypeDir
		if !directory && header.Typeflag != tar.TypeReg && header.Typeflag != tar.TypeRegA {
			_ = gzipReader.Close()
			return fmt.Errorf("%w: special tar entry", ErrUADashboardArchiveUnsafe)
		}
		if header.Size < 0 || header.Size > limits.FileBytes {
			_ = gzipReader.Close()
			return fmt.Errorf("%w: tar entry exceeds limit", ErrUADashboardArchiveLimit)
		}
		size := header.Size
		err = consume(uaArchiveEntry{Name: header.Name, Directory: directory, Read: func() ([]byte, error) {
			data, err := io.ReadAll(io.LimitReader(tarReader, limits.FileBytes+1))
			if err != nil || int64(len(data)) != size {
				return nil, fmt.Errorf("%w: tar entry size mismatch", ErrUADashboardArchiveUnsafe)
			}
			return data, nil
		}})
		if err != nil {
			_ = gzipReader.Close()
			return err
		}
	}
	if err := gzipReader.Close(); err != nil {
		return fmt.Errorf("%w: invalid gzip footer", ErrUADashboardArchiveUnsafe)
	}
	if _, err := buffered.Peek(1); !errors.Is(err, io.EOF) {
		return fmt.Errorf("%w: trailing gzip payload", ErrUADashboardArchiveUnsafe)
	}
	return nil
}

func normalizeUADashboardArchivePath(name string, directory bool) (string, error) {
	if name == "" || strings.ContainsAny(name, "\\\x00") || strings.Contains(name, "%") || strings.HasPrefix(name, "/") || strings.Contains(name, ":") {
		return "", fmt.Errorf("%w: unsafe archive path", ErrUADashboardArchiveUnsafe)
	}
	if directory {
		name = strings.TrimSuffix(name, "/")
	}
	if name == "" {
		return "", fmt.Errorf("%w: empty archive path", ErrUADashboardArchiveUnsafe)
	}
	for _, segment := range strings.Split(name, "/") {
		if segment == "" || segment == "." || segment == ".." {
			return "", fmt.Errorf("%w: ambiguous archive path", ErrUADashboardArchiveUnsafe)
		}
	}
	return name, nil
}

func uaDashboardArchiveFileLocation(path string) (string, string, bool) {
	if _, ok := uaDashboardAllowedFiles[path]; ok {
		return "root", path, true
	}
	if strings.HasPrefix(path, ".ua/") {
		basename := strings.TrimPrefix(path, ".ua/")
		if _, ok := uaDashboardAllowedFiles[basename]; ok {
			return "ua", basename, true
		}
	}
	return "", "", false
}

func validateStagedUADashboard(files map[string][]byte, limits UADashboardImportLimits) (UADashboardManifest, error) {
	paths := make([]string, 0, len(files))
	for path, data := range files {
		if err := validateUADashboardJSON(data, limits); err != nil {
			return UADashboardManifest{}, err
		}
		paths = append(paths, path)
	}
	sort.Strings(paths)
	manifest := UADashboardManifest{SchemaVersion: 1, Files: make([]UADashboardManifestFile, 0, len(paths))}
	for _, path := range paths {
		manifest.Files = append(manifest.Files, UADashboardManifestFile{Path: path, Bytes: int64(len(files[path]))})
	}
	return manifest, nil
}