package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/skip2/go-qrcode"
)

// ScanResult holds the result of a GoTestWAF scan
type ScanResult struct {
	EmailDomain string
	ReportFile  string
	Status      string // "pending", "running", "completed", "failed"
}

var (
	// in-memory storage for scan results
	scanResults = make(map[string]*ScanResult)
	mu          sync.RWMutex
	// emailRegex validates basic email format
	emailRegex = regexp.MustCompile(`^[^@]+@[^@]+\.[^@]+$`)
)

func main() {
	// Create main context for graceful shutdown
	mainCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	router := gin.Default()
	router.LoadHTMLGlob("src/templates/*")

	// Serve static files (frontend)
	router.Static("/static", "./src/static")
	router.StaticFile("/", "./src/static/index.html")
	router.StaticFile("/simulation", "./src/static/simulation.html")
	// Serve GoTestWAF reports
	router.Static("/reports", "./reports")

	// QR code generator endpoint
	router.GET("/qrcode", func(c *gin.Context) {
		text := c.Query("text")
		if text == "" {
			text = "WX-2026"
		}
		png, err := qrcode.Encode(text, qrcode.Medium, 256)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate QR code"})
			return
		}
		c.Data(http.StatusOK, "image/png", png)
	})

	// API endpoint: start security scan
	router.POST("/start", func(c *gin.Context) {
		email := strings.TrimSpace(c.PostForm("email"))
		if email == "" {
			c.JSON(http.StatusBadRequest, gin.H{"error": "Email is required"})
			return
		}

		// Validate email format
		if !emailRegex.MatchString(email) {
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid email format"})
			return
		}

		// Extract domain from email
		parts := strings.Split(email, "@")
		domain := parts[1] // safe because regex ensures exactly one @

		// Domain validation
		if !strings.Contains(domain, ".") {
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid domain"})
			return
		}

		// Prevent path traversal in domain (no slashes or backslashes)
		if strings.ContainsAny(domain, "/\\") {
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid domain characters"})
			return
		}

		// Limit domain length
		if len(domain) > 255 {
			c.JSON(http.StatusBadRequest, gin.H{"error": "Domain too long"})
			return
		}

		// Store scan result (pending)
		mu.Lock()
		scanResults[domain] = &ScanResult{
			EmailDomain: domain,
			ReportFile:  fmt.Sprintf("report_%s.html", strings.ReplaceAll(domain, ".", "_")),
			Status:      "pending",
		}
		mu.Unlock()

		// Trigger background scan (goroutine)
		go startGoTestWAFScan(mainCtx, domain)

		// Redirect to simulation page with domain as query param
		c.Redirect(http.StatusFound, "/simulation?domain="+domain)
	})

	// Admin dashboard: list all scan results
	router.GET("/admin-dashboard", func(c *gin.Context) {
		mu.RLock()
		defer mu.RUnlock()

		results := make([]ScanResult, 0, len(scanResults))
		completed := 0
		running := 0
		for _, r := range scanResults {
			results = append(results, *r)
			switch r.Status {
			case "completed":
				completed++
			case "running":
				running++
			}
		}

		c.HTML(http.StatusOK, "admin.html", gin.H{
			"Results":   results,
			"Total":     len(results),
			"Completed": completed,
			"Running":   running,
		})
	})

	// Endpoint to check scan status (for frontend polling)
	router.GET("/scan-status/:domain", func(c *gin.Context) {
		domain := c.Param("domain")
		mu.RLock()
		result, exists := scanResults[domain]
		mu.RUnlock()

		if !exists {
			c.JSON(http.StatusNotFound, gin.H{"error": "Domain not found"})
			return
		}

		c.JSON(http.StatusOK, gin.H{
			"domain": result.EmailDomain,
			"status": result.Status,
			"report": result.ReportFile,
		})
	})

	// Handle OS signals for graceful shutdown
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)

	// Run server in goroutine
	srv := &http.Server{
		Addr:    "0.0.0.0:8080",
		Handler: router,
	}

	go func() {
		fmt.Printf("Server starting on %s\n", srv.Addr)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			fmt.Printf("Server error: %v\n", err)
		}
	}()

	// Wait for shutdown signal
	<-signalChan
	fmt.Println("Shutdown signal received, shutting down gracefully...")

	// Cancel main context (signals all goroutines)
	cancel()

	// Create shutdown context with timeout
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer shutdownCancel()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		fmt.Printf("Server shutdown error: %v\n", err)
	}

	fmt.Println("Server stopped")
}

// startGoTestWAFScan launches actual GoTestWAF binary in background
func startGoTestWAFScan(parentCtx context.Context, domain string) {
	mu.Lock()
	if result, exists := scanResults[domain]; exists {
		result.Status = "running"
	}
	mu.Unlock()

	// Determine project root (AttackSurface/) relative to current working directory
	// Assuming the Go program is run from AttackSurface directory (where bin/ and reports/ are)
	projectRoot := "."
	reportsDir := filepath.Join(projectRoot, "reports")
	if err := os.MkdirAll(reportsDir, 0755); err != nil {
		fmt.Printf("Failed to create reports directory: %v\n", err)
		updateScanStatus(domain, "failed")
		return
	}

	// Prepare GoTestWAF command
	// Use the actual target server with Wallarm filtering node in monitoring mode
	targetURL := "https://git.sechpoint.app"
	binaryPath := filepath.Join(projectRoot, "bin", "gotestwaf")

	// Verify binary exists
	if _, err := os.Stat(binaryPath); os.IsNotExist(err) {
		fmt.Printf("GoTestWAF binary not found at %s\n", binaryPath)
		updateScanStatus(domain, "failed")
		return
	}

	// Create context with 120-second timeout for GoTestWAF scan
	scanCtx, cancel := context.WithTimeout(parentCtx, 120*time.Second)
	defer cancel()

	cmd := exec.CommandContext(scanCtx, binaryPath,
		"--url", targetURL,
		"--configPath", "src/gotestwaf/config.yaml",
		"--maxIdleConns", "2",
		"--maxRedirects", "10",
		"--reportPath", reportsDir,
		"--reportName", fmt.Sprintf("report_%s", strings.ReplaceAll(domain, ".", "_")),
		"--reportFormat", "html",
		"--wafName", "generic",
		"--skipWAFBlockCheck",
		"--nonBlockedAsPassed",
		"--tlsVerify",
		"--noEmailReport",
		"--quiet",
	)

	// Set working directory to project root (AttackSurface/)
	cmd.Dir = projectRoot

	// Run command in goroutine
	go func() {
		fmt.Printf("Starting GoTestWAF scan for domain: %s (target: %s)\n", domain, targetURL)
		startTime := time.Now()

		output, err := cmd.CombinedOutput()
		if err != nil {
			// Check if error is due to timeout
			if scanCtx.Err() == context.DeadlineExceeded {
				fmt.Printf("GoTestWAF scan timed out after 120s for domain: %s\n", domain)
				updateScanStatus(domain, "failed")
				return
			}
			fmt.Printf("GoTestWAF scan failed for %s: %v\nOutput:\n%s\n", domain, err, output)
			updateScanStatus(domain, "failed")
			return
		}

		elapsed := time.Since(startTime)
		fmt.Printf("GoTestWAF scan completed for %s in %v\n", domain, elapsed)
		updateScanStatus(domain, "completed")
	}()
}

// updateScanStatus updates the scan result status
func updateScanStatus(domain string, status string) {
	mu.Lock()
	defer mu.Unlock()
	if result, exists := scanResults[domain]; exists {
		result.Status = status
	}
}