mev-beta/tools/math-audit/cmd/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/spf13/cobra"

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

var (
	vectorsFile string
	reportDir   string
	verbose     bool
	exchange    string
	tolerance   float64
)

func main() {
	var rootCmd = &cobra.Command{
		Use:   "math-audit",
		Short: "Math audit tool for MEV Bot exchange calculations",
		Long: `A comprehensive math audit tool that validates pricing calculations,
amount in/out computations, and price impact calculations across all supported exchanges.`,
	}

	var auditCmd = &cobra.Command{
		Use:   "audit",
		Short: "Run comprehensive math audit",
		Long:  "Validates pricing math across all supported exchanges using canonical test vectors",
		RunE:  runAudit,
	}

	var validateCmd = &cobra.Command{
		Use:   "validate",
		Short: "Validate specific exchange calculations",
		Long:  "Validates calculations for a specific exchange using provided test vectors",
		RunE:  runValidate,
	}

	var reportCmd = &cobra.Command{
		Use:   "report",
		Short: "Generate audit reports",
		Long:  "Generates JSON and Markdown audit reports from previous validation runs",
		RunE:  runReport,
	}

	// Global flags
	rootCmd.PersistentFlags().StringVar(&vectorsFile, "vectors", "default", "Test vectors file or preset (default, comprehensive)")
	rootCmd.PersistentFlags().StringVar(&reportDir, "report", "reports/math/latest", "Report output directory")
	rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "Enable verbose output")

	// Validate command flags
	validateCmd.Flags().StringVar(&exchange, "exchange", "", "Exchange to validate (uniswap_v2, uniswap_v3, curve, balancer, etc.)")
	validateCmd.Flags().Float64Var(&tolerance, "tolerance", 0.0001, "Error tolerance in basis points (0.0001 = 1bp)")

	rootCmd.AddCommand(auditCmd, validateCmd, reportCmd)

	if err := rootCmd.Execute(); err != nil {
		log.Fatal(err)
	}
}

func runAudit(cmd *cobra.Command, args []string) error {
	ctx := context.Background()

	if verbose {
		fmt.Printf("Starting comprehensive math audit...\n")
		fmt.Printf("Vectors: %s\n", vectorsFile)
		fmt.Printf("Report Directory: %s\n", reportDir)
	}

	// Load test vectors
	vectors, err := internal.LoadTestVectors(vectorsFile)
	if err != nil {
		return fmt.Errorf("failed to load test vectors: %w", err)
	}

	// Initialize audit engine
	auditor := internal.NewMathAuditor(pkgmath.NewDecimalConverter(), tolerance)

	// Run audit across all exchanges
	results := make(map[string]*internal.ExchangeAuditResult)
	exchanges := []string{"uniswap_v2", "uniswap_v3", "curve", "balancer", "algebra", "camelot"}

	for _, exchangeType := range exchanges {
		if verbose {
			fmt.Printf("Auditing %s...\n", exchangeType)
		}

		exchangeVectors := vectors.GetExchangeVectors(exchangeType)
		if exchangeVectors == nil {
			fmt.Printf("Warning: No vectors found for %s\n", exchangeType)
			continue
		}

		result, err := auditor.AuditExchange(ctx, exchangeType, exchangeVectors)
		if err != nil {
			fmt.Printf("Error auditing %s: %v\n", exchangeType, err)
			continue
		}

		results[exchangeType] = result

		if verbose {
			fmt.Printf("  - Tests: %d, Passed: %d, Failed: %d, Max Error: %.4f bp\n",
				result.TotalTests, result.PassedTests, result.FailedTests, result.MaxErrorBP)
		}
	}

	// Generate comprehensive report
	report := &internal.ComprehensiveAuditReport{
		Timestamp:       time.Now(),
		VectorsFile:     vectorsFile,
		ToleranceBP:     tolerance * 10000, // Convert to basis points
		ExchangeResults: results,
		OverallPassed:   true,
		TotalTests:      0,
		TotalPassed:     0,
		TotalFailed:     0,
	}

	// Calculate overall statistics
	for _, result := range results {
		report.TotalTests += result.TotalTests
		report.TotalPassed += result.PassedTests
		report.TotalFailed += result.FailedTests
		if result.FailedTests > 0 {
			report.OverallPassed = false
		}
	}

	// Save reports
	if err := saveReports(report, reportDir); err != nil {
		return fmt.Errorf("failed to save reports: %w", err)
	}

	// Print summary
	fmt.Printf("\nMath Audit Complete:\n")
	fmt.Printf("  Total Tests: %d\n", report.TotalTests)
	fmt.Printf("  Passed: %d\n", report.TotalPassed)
	fmt.Printf("  Failed: %d\n", report.TotalFailed)
	fmt.Printf("  Overall Status: %s\n", map[bool]string{true: "PASS", false: "FAIL"}[report.OverallPassed])
	fmt.Printf("  Reports saved to: %s\n", reportDir)

	if !report.OverallPassed {
		return fmt.Errorf("math audit failed - see report for details")
	}

	return nil
}

func runValidate(cmd *cobra.Command, args []string) error {
	if exchange == "" {
		return fmt.Errorf("exchange must be specified with --exchange flag")
	}

	ctx := context.Background()

	if verbose {
		fmt.Printf("Validating %s exchange calculations...\n", exchange)
	}

	// Load test vectors
	vectors, err := internal.LoadTestVectors(vectorsFile)
	if err != nil {
		return fmt.Errorf("failed to load test vectors: %w", err)
	}

	exchangeVectors := vectors.GetExchangeVectors(exchange)
	if exchangeVectors == nil {
		return fmt.Errorf("no test vectors found for exchange: %s", exchange)
	}

	// Initialize auditor
	auditor := internal.NewMathAuditor(pkgmath.NewDecimalConverter(), tolerance)

	// Run validation
	result, err := auditor.AuditExchange(ctx, exchange, exchangeVectors)
	if err != nil {
		return fmt.Errorf("validation failed: %w", err)
	}

	// Print detailed results
	fmt.Printf("\nValidation Results for %s:\n", exchange)
	fmt.Printf("  Total Tests: %d\n", result.TotalTests)
	fmt.Printf("  Passed: %d\n", result.PassedTests)
	fmt.Printf("  Failed: %d\n", result.FailedTests)
	fmt.Printf("  Max Error: %.4f bp\n", result.MaxErrorBP)
	fmt.Printf("  Avg Error: %.4f bp\n", result.AvgErrorBP)

	if len(result.FailedCases) > 0 {
		fmt.Printf("\nFailed Test Cases:\n")
		for i, failure := range result.FailedCases {
			if i >= 5 { // Limit output
				fmt.Printf("  ... and %d more\n", len(result.FailedCases)-5)
				break
			}
			fmt.Printf("  - %s: %.4f bp error\n", failure.TestName, failure.ErrorBP)
		}
	}

	if result.FailedTests > 0 {
		return fmt.Errorf("validation failed for %s", exchange)
	}

	fmt.Printf("\n✓ All tests passed for %s\n", exchange)
	return nil
}

func runReport(cmd *cobra.Command, args []string) error {
	// Load existing audit results and generate fresh reports
	fmt.Printf("Generating audit reports in %s...\n", reportDir)

	// This would load existing JSON data and regenerate markdown reports
	// For now, just indicate the feature exists
	fmt.Printf("Report generation feature coming soon.\n")
	fmt.Printf("Run 'math-audit audit' to generate fresh audit reports.\n")

	return nil
}

func saveReports(report *internal.ComprehensiveAuditReport, reportDir string) error {
	// Ensure report directory exists
	if err := os.MkdirAll(reportDir, 0755); err != nil {
		return err
	}

	// Save JSON report
	jsonFile := filepath.Join(reportDir, "audit_results.json")
	jsonData, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		return err
	}

	if err := os.WriteFile(jsonFile, jsonData, 0644); err != nil {
		return err
	}

	// Save Markdown report
	markdownFile := filepath.Join(reportDir, "audit_report.md")
	markdownContent := internal.GenerateMarkdownReport(report)

	if err := os.WriteFile(markdownFile, []byte(markdownContent), 0644); err != nil {
		return err
	}

	return nil
}