package metrics

import (
	"sort"
	"sync"
	"time"
)

// Collector handles metrics collection during stress testing
type Collector struct {
	mu              sync.RWMutex
	responseTimes   []time.Duration
	statusCodes     map[int]int64
	errors          []error
	totalRequests   int64
	successRequests int64
	failedRequests  int64
	startTime       time.Time
}

// Results contains aggregated test results
type Results struct {
	TotalRequests     int64
	SuccessRequests   int64
	FailedRequests    int64
	MinResponseTime   time.Duration
	MaxResponseTime   time.Duration
	AvgResponseTime   time.Duration
	P50               time.Duration
	P90               time.Duration
	P95               time.Duration
	P99               time.Duration
	StatusCodes       map[int]int64
	Errors            []error
	TestDuration      time.Duration
	RequestsPerSecond float64
}

// NewCollector creates a new metrics collector
func NewCollector() *Collector {
	return &Collector{
		responseTimes: make([]time.Duration, 0),
		statusCodes:   make(map[int]int64),
		errors:        make([]error, 0),
		startTime:     time.Now(),
	}
}

// RecordRequest records a single request's metrics
func (c *Collector) RecordRequest(responseTime time.Duration, statusCode int, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.totalRequests++
	c.responseTimes = append(c.responseTimes, responseTime)

	if err != nil {
		c.failedRequests++
		c.errors = append(c.errors, err)
	} else {
		c.successRequests++
		c.statusCodes[statusCode]++
	}
}

// GetResults calculates and returns aggregated results
func (c *Collector) GetResults() *Results {
	c.mu.RLock()
	defer c.mu.RUnlock()

	testDuration := time.Since(c.startTime)
	rps := float64(c.totalRequests) / testDuration.Seconds()

	if len(c.responseTimes) == 0 {
		return &Results{
			TotalRequests:     c.totalRequests,
			SuccessRequests:   c.successRequests,
			FailedRequests:    c.failedRequests,
			StatusCodes:       copyMap(c.statusCodes),
			Errors:            copyErrors(c.errors),
			TestDuration:      testDuration,
			RequestsPerSecond: rps,
		}
	}

	// Filter out zero response times (from errors) and sort for percentile calculations
	var validTimes []time.Duration
	for _, t := range c.responseTimes {
		if t > 0 {
			validTimes = append(validTimes, t)
		}
	}

	if len(validTimes) == 0 {
		return &Results{
			TotalRequests:     c.totalRequests,
			SuccessRequests:   c.successRequests,
			FailedRequests:    c.failedRequests,
			StatusCodes:       copyMap(c.statusCodes),
			Errors:            copyErrors(c.errors),
			TestDuration:      testDuration,
			RequestsPerSecond: rps,
		}
	}

	sort.Slice(validTimes, func(i, j int) bool {
		return validTimes[i] < validTimes[j]
	})

	// Calculate statistics
	minTime := validTimes[0]
	maxTime := validTimes[len(validTimes)-1]
	avgTime := calculateAverage(validTimes)

	return &Results{
		TotalRequests:     c.totalRequests,
		SuccessRequests:   c.successRequests,
		FailedRequests:    c.failedRequests,
		MinResponseTime:   minTime,
		MaxResponseTime:   maxTime,
		AvgResponseTime:   avgTime,
		P50:               calculatePercentile(validTimes, 50),
		P90:               calculatePercentile(validTimes, 90),
		P95:               calculatePercentile(validTimes, 95),
		P99:               calculatePercentile(validTimes, 99),
		StatusCodes:       copyMap(c.statusCodes),
		Errors:            copyErrors(c.errors),
		TestDuration:      testDuration,
		RequestsPerSecond: rps,
	}
}

// Reset clears all collected metrics
func (c *Collector) Reset() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.responseTimes = c.responseTimes[:0]
	c.statusCodes = make(map[int]int64)
	c.errors = c.errors[:0]
	c.totalRequests = 0
	c.successRequests = 0
	c.failedRequests = 0
	c.startTime = time.Now()
}

// calculatePercentile calculates the nth percentile of sorted response times
func calculatePercentile(sortedTimes []time.Duration, percentile int) time.Duration {
	if len(sortedTimes) == 0 {
		return 0
	}

	index := int(float64(len(sortedTimes)-1) * float64(percentile) / 100.0)
	if index >= len(sortedTimes) {
		index = len(sortedTimes) - 1
	}

	return sortedTimes[index]
}

// calculateAverage calculates the average response time
func calculateAverage(times []time.Duration) time.Duration {
	if len(times) == 0 {
		return 0
	}

	var total time.Duration
	for _, t := range times {
		total += t
	}

	return total / time.Duration(len(times))
}

// copyMap creates a copy of the status codes map
func copyMap(original map[int]int64) map[int]int64 {
	copy := make(map[int]int64)
	for k, v := range original {
		copy[k] = v
	}
	return copy
}

// copyErrors creates a copy of the errors slice
func copyErrors(original []error) []error {
	copy := make([]error, len(original))
	copy = append(copy, original...)
	return copy
}