package loader

import (
	"encoding/json"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	"github.com/fraktal/mev-beta/tools/math-audit/internal/models"
)

const defaultVectorDir = "tools/math-audit/vectors"

// LoadVectors loads vectors based on the selector. The selector can be:
//   - "default" (load all embedded vectors)
//   - a comma-separated list of files or directories.
func LoadVectors(selector string) ([]models.Vector, error) {
	if selector == "" || selector == "default" {
		return loadFromDir(defaultVectorDir)
	}

	var all []models.Vector
	parts := strings.Split(selector, ",")
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}

		info, err := os.Stat(part)
		if err != nil {
			return nil, fmt.Errorf("stat %s: %w", part, err)
		}

		if info.IsDir() {
			vecs, err := loadFromDir(part)
			if err != nil {
				return nil, err
			}
			all = append(all, vecs...)
			continue
		}

		vec, err := loadFromFile(part)
		if err != nil {
			return nil, err
		}
		all = append(all, vec)
	}

	return all, nil
}

func loadFromDir(dir string) ([]models.Vector, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, fmt.Errorf("read dir %s: %w", dir, err)
	}

	var vectors []models.Vector
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		if !strings.HasSuffix(entry.Name(), ".json") {
			continue
		}

		path := filepath.Join(dir, entry.Name())
		vec, err := loadFromFile(path)
		if err != nil {
			return nil, err
		}
		vectors = append(vectors, vec)
	}

	return vectors, nil
}

func loadFromFile(path string) (models.Vector, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return models.Vector{}, fmt.Errorf("read vector %s: %w", path, err)
	}

	var vec models.Vector
	if err := json.Unmarshal(data, &vec); err != nil {
		return models.Vector{}, fmt.Errorf("decode vector %s: %w", path, err)
	}

	if err := vec.Validate(); err != nil {
		return models.Vector{}, fmt.Errorf("validate vector %s: %w", path, err)
	}

	return vec, nil
}

// WalkVectors walks every vector JSON file in a directory hierarchy.
func WalkVectors(root string, fn func(path string, vec models.Vector) error) error {
	return filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() || !strings.HasSuffix(d.Name(), ".json") {
			return nil
		}
		vec, err := loadFromFile(path)
		if err != nil {
			return err
		}
		return fn(path, vec)
	})
}