package format

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
)

var (
	importRe       = regexp.MustCompile(`(?ms)^import \(([^)]+)\)`)
	singleImportRe = regexp.MustCompile(`(?ms)^import "[^"]+"`)
	otherRe        = regexp.MustCompile(`^(?:var|const|func)\s`)
	ErrNoImports   = errors.New("no imports found")
)

// Source formats a given src's imports
func Source(src []byte, module string) ([]byte, error) {
	if singleImportRe.Match(src) {
		return src, nil
	}

	importStart := importRe.FindIndex(src)
	if importStart == nil {
		return nil, fmt.Errorf("could not find imports: %w", ErrNoImports)
	}

	otherStart := otherRe.FindIndex(src)
	if otherStart != nil && otherStart[0] < importStart[0] {
		return nil, fmt.Errorf("found non-imports before imports: %w", ErrNoImports)
	}

	groups := importRe.FindStringSubmatch(string(src))
	if groups[0] == "" {
		return nil, fmt.Errorf("could not find imports: %w", ErrNoImports)
	}

	imports := strings.Split(groups[1], "\n")
	for idx, i := range imports {
		imports[idx] = strings.TrimSpace(i)
	}

	block := parseBlock(module, imports)
	replaced := strings.Replace(string(src), groups[0], block.String(), 1)
	return []byte(replaced), nil
}