package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

// Configuration
type LoadTestConfig struct {
	GatewayURL      string        `json:"gateway_url"`
	ConcurrentUsers int           `json:"concurrent_users"`
	TestDuration    time.Duration `json:"test_duration"`
	FileSize        int64         `json:"file_size"`
	RampUpTime      time.Duration `json:"ramp_up_time"`
	ReportInterval  time.Duration `json:"report_interval"`
}

// Metrics
type Metrics struct {
	TotalRequests       int64         `json:"total_requests"`
	SuccessfulRequests  int64         `json:"successful_requests"`
	FailedRequests      int64         `json:"failed_requests"`
	TotalBytesUploaded  int64         `json:"total_bytes_uploaded"`
	TotalBytesDownloaded int64        `json:"total_bytes_downloaded"`
	AverageResponseTime time.Duration `json:"average_response_time"`
	MinResponseTime     time.Duration `json:"min_response_time"`
	MaxResponseTime     time.Duration `json:"max_response_time"`
	RequestsPerSecond   float64       `json:"requests_per_second"`
	BytesPerSecond      float64       `json:"bytes_per_second"`
	ErrorRate           float64       `json:"error_rate"`
	P95ResponseTime     time.Duration `json:"p95_response_time"`
	P99ResponseTime     time.Duration `json:"p99_response_time"`
}

// Request result
type RequestResult struct {
	Success      bool
	ResponseTime time.Duration
	BytesTransferred int64
	ErrorMessage string
	RequestType  string
}

// LoadTester manages the load testing process
type LoadTester struct {
	config        LoadTestConfig
	httpClient    *http.Client
	metrics       *Metrics
	responseTimes []time.Duration
	mu            sync.RWMutex
	ctx           context.Context
	cancel        context.CancelFunc
}

// NewLoadTester creates a new load tester instance
func NewLoadTester(config LoadTestConfig) *LoadTester {
	ctx, cancel := context.WithCancel(context.Background())
	
	return &LoadTester{
		config: config,
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
			Transport: &http.Transport{
				MaxIdleConns:        100,
				MaxIdleConnsPerHost: 10,
				IdleConnTimeout:     90 * time.Second,
			},
		},
		metrics:       &Metrics{},
		responseTimes: make([]time.Duration, 0, 10000),
		ctx:           ctx,
		cancel:        cancel,
	}
}

// generateTestData creates random test data
func (lt *LoadTester) generateTestData(size int64) []byte {
	data := make([]byte, size)
	if _, err := rand.Read(data); err != nil {
		log.Printf("Failed to generate random data: %v", err)
		// Fallback to pattern-based data
		for i := range data {
			data[i] = byte(i % 256)
		}
	}
	return data
}

// uploadFile simulates file upload
func (lt *LoadTester) uploadFile(workerID int, fileData []byte) RequestResult {
	start := time.Now()
	
	// Create multipart form
	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)
	
	// Create file field
	fileWriter, err := writer.CreateFormFile("file", fmt.Sprintf("load_test_%d_%d.bin", workerID, time.Now().UnixNano()))
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: time.Since(start),
			ErrorMessage: fmt.Sprintf("Failed to create form file: %v", err),
			RequestType:  "upload",
		}
	}
	
	if _, err := fileWriter.Write(fileData); err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: time.Since(start),
			ErrorMessage: fmt.Sprintf("Failed to write file data: %v", err),
			RequestType:  "upload",
		}
	}
	
	writer.Close()
	
	// Create request
	req, err := http.NewRequestWithContext(lt.ctx, "POST", lt.config.GatewayURL+"/upload", &buf)
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: time.Since(start),
			ErrorMessage: fmt.Sprintf("Failed to create request: %v", err),
			RequestType:  "upload",
		}
	}
	
	req.Header.Set("Content-Type", writer.FormDataContentType())
	
	// Add test authentication header
	testPubkey := "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
	sessionToken := "test_session_token_" + testPubkey
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	
	// Send request
	resp, err := lt.httpClient.Do(req)
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: time.Since(start),
			ErrorMessage: fmt.Sprintf("Request failed: %v", err),
			RequestType:  "upload",
		}
	}
	defer resp.Body.Close()
	
	responseTime := time.Since(start)
	
	// Read response
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: responseTime,
			ErrorMessage: fmt.Sprintf("Failed to read response: %v", err),
			RequestType:  "upload",
		}
	}
	
	if resp.StatusCode != http.StatusOK {
		return RequestResult{
			Success:      false,
			ResponseTime: responseTime,
			ErrorMessage: fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(respBody)),
			RequestType:  "upload",
		}
	}
	
	// Parse response to get file hash for potential download test
	var uploadResp map[string]interface{}
	if err := json.Unmarshal(respBody, &uploadResp); err != nil {
		log.Printf("Warning: Failed to parse upload response: %v", err)
	}
	
	return RequestResult{
		Success:          true,
		ResponseTime:     responseTime,
		BytesTransferred: int64(len(fileData)),
		RequestType:      "upload",
	}
}

// downloadFile simulates file download
func (lt *LoadTester) downloadFile(fileHash string) RequestResult {
	start := time.Now()
	
	req, err := http.NewRequestWithContext(lt.ctx, "GET", lt.config.GatewayURL+"/download/"+fileHash, nil)
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: time.Since(start),
			ErrorMessage: fmt.Sprintf("Failed to create request: %v", err),
			RequestType:  "download",
		}
	}
	
	resp, err := lt.httpClient.Do(req)
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: time.Since(start),
			ErrorMessage: fmt.Sprintf("Request failed: %v", err),
			RequestType:  "download",
		}
	}
	defer resp.Body.Close()
	
	responseTime := time.Since(start)
	
	// Read response body to measure bytes transferred
	bytesRead, err := io.Copy(io.Discard, resp.Body)
	if err != nil {
		return RequestResult{
			Success:      false,
			ResponseTime: responseTime,
			ErrorMessage: fmt.Sprintf("Failed to read response: %v", err),
			RequestType:  "download",
		}
	}
	
	if resp.StatusCode != http.StatusOK {
		return RequestResult{
			Success:      false,
			ResponseTime: responseTime,
			ErrorMessage: fmt.Sprintf("HTTP %d", resp.StatusCode),
			RequestType:  "download",
		}
	}
	
	return RequestResult{
		Success:          true,
		ResponseTime:     responseTime,
		BytesTransferred: bytesRead,
		RequestType:      "download",
	}
}

// worker simulates a concurrent user
func (lt *LoadTester) worker(workerID int, results chan<- RequestResult, wg *sync.WaitGroup) {
	defer wg.Done()
	
	fileData := lt.generateTestData(lt.config.FileSize)
	
	for {
		select {
		case <-lt.ctx.Done():
			return
		default:
			// Perform upload test
			result := lt.uploadFile(workerID, fileData)
			results <- result
			
			// Small delay between requests to prevent overwhelming
			time.Sleep(time.Millisecond * 100)
		}
	}
}

// updateMetrics updates the metrics with new result
func (lt *LoadTester) updateMetrics(result RequestResult) {
	lt.mu.Lock()
	defer lt.mu.Unlock()
	
	atomic.AddInt64(&lt.metrics.TotalRequests, 1)
	
	if result.Success {
		atomic.AddInt64(&lt.metrics.SuccessfulRequests, 1)
		if result.RequestType == "upload" {
			atomic.AddInt64(&lt.metrics.TotalBytesUploaded, result.BytesTransferred)
		} else {
			atomic.AddInt64(&lt.metrics.TotalBytesDownloaded, result.BytesTransferred)
		}
	} else {
		atomic.AddInt64(&lt.metrics.FailedRequests, 1)
		if result.ErrorMessage != "" {
			log.Printf("Request failed: %s", result.ErrorMessage)
		}
	}
	
	// Track response times
	lt.responseTimes = append(lt.responseTimes, result.ResponseTime)
	
	// Update min/max response times
	if lt.metrics.MinResponseTime == 0 || result.ResponseTime < lt.metrics.MinResponseTime {
		lt.metrics.MinResponseTime = result.ResponseTime
	}
	if result.ResponseTime > lt.metrics.MaxResponseTime {
		lt.metrics.MaxResponseTime = result.ResponseTime
	}
}

// calculateStatistics computes statistical metrics
func (lt *LoadTester) calculateStatistics() {
	lt.mu.Lock()
	defer lt.mu.Unlock()
	
	if len(lt.responseTimes) == 0 {
		return
	}
	
	// Calculate average response time
	var totalResponseTime time.Duration
	for _, rt := range lt.responseTimes {
		totalResponseTime += rt
	}
	lt.metrics.AverageResponseTime = totalResponseTime / time.Duration(len(lt.responseTimes))
	
	// Sort response times for percentile calculations
	responseTimes := make([]time.Duration, len(lt.responseTimes))
	copy(responseTimes, lt.responseTimes)
	
	// Simple sort (for small datasets)
	for i := 0; i < len(responseTimes)-1; i++ {
		for j := i + 1; j < len(responseTimes); j++ {
			if responseTimes[i] > responseTimes[j] {
				responseTimes[i], responseTimes[j] = responseTimes[j], responseTimes[i]
			}
		}
	}
	
	// Calculate percentiles
	if len(responseTimes) > 0 {
		p95Index := int(float64(len(responseTimes)) * 0.95)
		p99Index := int(float64(len(responseTimes)) * 0.99)
		
		if p95Index >= len(responseTimes) {
			p95Index = len(responseTimes) - 1
		}
		if p99Index >= len(responseTimes) {
			p99Index = len(responseTimes) - 1
		}
		
		lt.metrics.P95ResponseTime = responseTimes[p95Index]
		lt.metrics.P99ResponseTime = responseTimes[p99Index]
	}
	
	// Calculate error rate
	if lt.metrics.TotalRequests > 0 {
		lt.metrics.ErrorRate = float64(lt.metrics.FailedRequests) / float64(lt.metrics.TotalRequests) * 100
	}
}

// printReport prints current performance metrics
func (lt *LoadTester) printReport(elapsed time.Duration) {
	lt.calculateStatistics()
	
	totalRequests := atomic.LoadInt64(&lt.metrics.TotalRequests)
	successfulRequests := atomic.LoadInt64(&lt.metrics.SuccessfulRequests)
	failedRequests := atomic.LoadInt64(&lt.metrics.FailedRequests)
	totalBytesUploaded := atomic.LoadInt64(&lt.metrics.TotalBytesUploaded)
	
	if elapsed.Seconds() > 0 {
		lt.metrics.RequestsPerSecond = float64(totalRequests) / elapsed.Seconds()
		lt.metrics.BytesPerSecond = float64(totalBytesUploaded) / elapsed.Seconds()
	}
	
	fmt.Printf("\n📊 Load Test Report (Elapsed: %v)\n", elapsed.Round(time.Second))
	fmt.Printf("====================================\n")
	fmt.Printf("Total Requests:     %d\n", totalRequests)
	fmt.Printf("Successful:         %d (%.1f%%)\n", successfulRequests, float64(successfulRequests)/float64(totalRequests)*100)
	fmt.Printf("Failed:             %d (%.1f%%)\n", failedRequests, lt.metrics.ErrorRate)
	fmt.Printf("Requests/sec:       %.2f\n", lt.metrics.RequestsPerSecond)
	fmt.Printf("Data Uploaded:      %.2f MB\n", float64(totalBytesUploaded)/(1024*1024))
	fmt.Printf("Upload Speed:       %.2f MB/s\n", lt.metrics.BytesPerSecond/(1024*1024))
	fmt.Printf("\nResponse Times:\n")
	fmt.Printf("  Average:          %v\n", lt.metrics.AverageResponseTime.Round(time.Millisecond))
	fmt.Printf("  Min:              %v\n", lt.metrics.MinResponseTime.Round(time.Millisecond))
	fmt.Printf("  Max:              %v\n", lt.metrics.MaxResponseTime.Round(time.Millisecond))
	fmt.Printf("  95th percentile:  %v\n", lt.metrics.P95ResponseTime.Round(time.Millisecond))
	fmt.Printf("  99th percentile:  %v\n", lt.metrics.P99ResponseTime.Round(time.Millisecond))
	
	// System resource usage
	var memStats runtime.MemStats
	runtime.ReadMemStats(&memStats)
	fmt.Printf("\nSystem Resources:\n")
	fmt.Printf("  Goroutines:       %d\n", runtime.NumGoroutine())
	fmt.Printf("  Memory Used:      %.2f MB\n", float64(memStats.Alloc)/(1024*1024))
	fmt.Printf("  Memory Total:     %.2f MB\n", float64(memStats.TotalAlloc)/(1024*1024))
	fmt.Printf("  GC Cycles:        %d\n", memStats.NumGC)
}

// saveResults saves detailed results to JSON file
func (lt *LoadTester) saveResults(filename string, testDuration time.Duration) error {
	lt.calculateStatistics()
	
	result := struct {
		Config      LoadTestConfig `json:"config"`
		Metrics     *Metrics       `json:"metrics"`
		TestInfo    map[string]interface{} `json:"test_info"`
	}{
		Config:  lt.config,
		Metrics: lt.metrics,
		TestInfo: map[string]interface{}{
			"test_duration":    testDuration.String(),
			"timestamp":        time.Now().Format(time.RFC3339),
			"go_version":       runtime.Version(),
			"num_cpu":          runtime.NumCPU(),
			"os":               runtime.GOOS,
			"arch":             runtime.GOARCH,
		},
	}
	
	data, err := json.MarshalIndent(result, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal results: %v", err)
	}
	
	return os.WriteFile(filename, data, 0644)
}

// Run executes the load test
func (lt *LoadTester) Run() error {
	fmt.Printf("🚀 Starting Load Test\n")
	fmt.Printf("=====================\n")
	fmt.Printf("Gateway URL:        %s\n", lt.config.GatewayURL)
	fmt.Printf("Concurrent Users:   %d\n", lt.config.ConcurrentUsers)
	fmt.Printf("Test Duration:      %v\n", lt.config.TestDuration)
	fmt.Printf("File Size:          %.2f MB\n", float64(lt.config.FileSize)/(1024*1024))
	fmt.Printf("Ramp Up Time:       %v\n", lt.config.RampUpTime)
	fmt.Printf("\n")
	
	// Test gateway connectivity
	fmt.Print("🔍 Testing gateway connectivity...")
	resp, err := lt.httpClient.Get(lt.config.GatewayURL + "/health")
	if err != nil {
		fmt.Printf(" ❌ FAILED\n")
		return fmt.Errorf("gateway not accessible: %v", err)
	}
	resp.Body.Close()
	fmt.Printf(" ✅ OK\n\n")
	
	// Start workers
	results := make(chan RequestResult, lt.config.ConcurrentUsers*2)
	var wg sync.WaitGroup
	
	startTime := time.Now()
	
	// Ramp up workers gradually
	for i := 0; i < lt.config.ConcurrentUsers; i++ {
		wg.Add(1)
		go lt.worker(i, results, &wg)
		
		// Stagger worker startup
		if lt.config.RampUpTime > 0 {
			time.Sleep(lt.config.RampUpTime / time.Duration(lt.config.ConcurrentUsers))
		}
	}
	
	// Results collector
	go func() {
		for result := range results {
			lt.updateMetrics(result)
		}
	}()
	
	// Report generator
	reportTicker := time.NewTicker(lt.config.ReportInterval)
	defer reportTicker.Stop()
	
	testTimer := time.NewTimer(lt.config.TestDuration)
	defer testTimer.Stop()
	
	fmt.Printf("🔥 Load test running... (Press Ctrl+C to stop early)\n")
	
	// Main test loop
	for {
		select {
		case <-testTimer.C:
			fmt.Printf("\n⏰ Test duration reached, stopping...\n")
			lt.cancel()
			goto finish
		case <-reportTicker.C:
			lt.printReport(time.Since(startTime))
		}
	}
	
finish:
	// Wait for workers to finish
	wg.Wait()
	close(results)
	
	// Wait a bit for final results to be processed
	time.Sleep(100 * time.Millisecond)
	
	testDuration := time.Since(startTime)
	
	// Final report
	fmt.Printf("\n🏁 Load Test Completed!\n")
	lt.printReport(testDuration)
	
	// Save results
	resultsFile := fmt.Sprintf("load_test_results_%s.json", time.Now().Format("20060102_150405"))
	if err := lt.saveResults(resultsFile, testDuration); err != nil {
		log.Printf("Failed to save results: %v", err)
	} else {
		fmt.Printf("\nResults saved to: %s\n", resultsFile)
	}
	
	// Performance recommendations
	lt.printRecommendations()
	
	return nil
}

// printRecommendations provides performance insights
func (lt *LoadTester) printRecommendations() {
	fmt.Printf("\n💡 Performance Insights:\n")
	fmt.Printf("========================\n")
	
	if lt.metrics.ErrorRate > 5 {
		fmt.Printf("⚠️  High error rate (%.1f%%) - consider reducing concurrent users or increasing server resources\n", lt.metrics.ErrorRate)
	}
	
	if lt.metrics.RequestsPerSecond < float64(lt.config.ConcurrentUsers)*0.1 {
		fmt.Printf("⚠️  Low throughput - potential bottlenecks in server or network\n")
	}
	
	if lt.metrics.P95ResponseTime > 5*time.Second {
		fmt.Printf("⚠️  High P95 response time (%v) - server may be under stress\n", lt.metrics.P95ResponseTime)
	}
	
	uploadSpeedMBps := lt.metrics.BytesPerSecond / (1024 * 1024)
	if uploadSpeedMBps > 100 {
		fmt.Printf("✅ Excellent upload performance (%.2f MB/s)\n", uploadSpeedMBps)
	} else if uploadSpeedMBps > 10 {
		fmt.Printf("✅ Good upload performance (%.2f MB/s)\n", uploadSpeedMBps)
	} else {
		fmt.Printf("⚠️  Upload performance could be improved (%.2f MB/s)\n", uploadSpeedMBps)
	}
	
	if lt.metrics.ErrorRate == 0 {
		fmt.Printf("✅ Perfect reliability - no failed requests\n")
	}
}

func main() {
	// Default configuration
	config := LoadTestConfig{
		GatewayURL:      "http://localhost:9876",
		ConcurrentUsers: 10,
		TestDuration:    2 * time.Minute,
		FileSize:        1024 * 1024, // 1MB
		RampUpTime:      10 * time.Second,
		ReportInterval:  15 * time.Second,
	}
	
	// Override with environment variables if present
	if url := os.Getenv("GATEWAY_URL"); url != "" {
		config.GatewayURL = url
	}
	
	if users := os.Getenv("CONCURRENT_USERS"); users != "" {
		if u, err := fmt.Sscanf(users, "%d", &config.ConcurrentUsers); err == nil && u == 1 {
			// Successfully parsed
		}
	}
	
	if duration := os.Getenv("TEST_DURATION"); duration != "" {
		if d, err := time.ParseDuration(duration); err == nil {
			config.TestDuration = d
		}
	}
	
	if size := os.Getenv("FILE_SIZE"); size != "" {
		if s, err := fmt.Sscanf(size, "%d", &config.FileSize); err == nil && s == 1 {
			// Successfully parsed
		}
	}
	
	// Create and run load tester
	tester := NewLoadTester(config)
	
	if err := tester.Run(); err != nil {
		log.Fatalf("Load test failed: %v", err)
	}
}