package proxy

import (
	"bufio"
	"context"
	"errors"
	"net"
	"net/http"
	"sync/atomic"
	"time"
)

// SlowDripWriter wraps http.ResponseWriter to stream bytes with delays.
// This simulates slow network connections or bandwidth throttling.
type SlowDripWriter struct {
	w           http.ResponseWriter
	bytesPerMs  int // Bytes to write per millisecond
	chunkSize   int // Size of each write chunk
	ctx         context.Context
	written     int64
	headersSent atomic.Bool
}

// NewSlowDripWriter creates a new slow-drip writer
func NewSlowDripWriter(w http.ResponseWriter, bytesPerMs, chunkSize int, ctx context.Context) *SlowDripWriter {
	if bytesPerMs <= 0 {
		bytesPerMs = 10 // Default: 10 bytes/ms = 10KB/s
	}
	if chunkSize <= 0 {
		chunkSize = 1 // Default: 1 byte chunks for maximum slowness
	}
	return &SlowDripWriter{
		w:          w,
		bytesPerMs: bytesPerMs,
		chunkSize:  chunkSize,
		ctx:        ctx,
	}
}

// Header returns the header map
func (sdw *SlowDripWriter) Header() http.Header {
	return sdw.w.Header()
}

// WriteHeader sends the HTTP response header with the provided status code
func (sdw *SlowDripWriter) WriteHeader(statusCode int) {
	if sdw.headersSent.CompareAndSwap(false, true) {
		sdw.w.WriteHeader(statusCode)
	}
}

// Write writes data with delays between chunks
func (sdw *SlowDripWriter) Write(p []byte) (int, error) {
	// Ensure headers are sent
	if !sdw.headersSent.Load() {
		sdw.WriteHeader(http.StatusOK)
	}

	written := 0
	for written < len(p) {
		// Check context cancellation
		select {
		case <-sdw.ctx.Done():
			return written, sdw.ctx.Err()
		default:
		}

		// Calculate chunk to write
		end := written + sdw.chunkSize
		if end > len(p) {
			end = len(p)
		}
		chunk := p[written:end]

		// Write chunk
		n, err := sdw.w.Write(chunk)
		written += n
		sdw.written += int64(n)

		if err != nil {
			return written, err
		}

		// Flush to ensure bytes are sent immediately
		if flusher, ok := sdw.w.(http.Flusher); ok {
			flusher.Flush()
		}

		// Calculate delay based on bytes written
		// bytesPerMs means we can write bytesPerMs bytes per millisecond
		// So for len(chunk) bytes, we wait (len(chunk) / bytesPerMs) milliseconds
		delayMs := len(chunk) / sdw.bytesPerMs
		if delayMs > 0 {
			select {
			case <-sdw.ctx.Done():
				return written, sdw.ctx.Err()
			case <-time.After(time.Duration(delayMs) * time.Millisecond):
			}
		}
	}

	return written, nil
}

// Flush implements http.Flusher
func (sdw *SlowDripWriter) Flush() {
	if flusher, ok := sdw.w.(http.Flusher); ok {
		flusher.Flush()
	}
}

// Hijack implements http.Hijacker
func (sdw *SlowDripWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := sdw.w.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, errors.New("underlying ResponseWriter does not support hijacking")
}

// BytesWritten returns the total bytes written
func (sdw *SlowDripWriter) BytesWritten() int64 {
	return sdw.written
}

// ConnectionDropWriter wraps http.ResponseWriter to drop connection mid-response.
// This simulates network failures, connection resets, or abrupt disconnections.
type ConnectionDropWriter struct {
	w                http.ResponseWriter
	dropAfterPercent float64 // Drop after this percentage of expected body
	dropAfterBytes   int64   // Drop after this many bytes (takes precedence)
	expectedSize     int64   // Expected total response size
	bytesWritten     int64
	dropped          atomic.Bool
	headersSent      atomic.Bool
}

// NewConnectionDropWriter creates a new connection drop writer
func NewConnectionDropWriter(w http.ResponseWriter, dropAfterPercent float64, dropAfterBytes int64, expectedSize int64) *ConnectionDropWriter {
	cdw := &ConnectionDropWriter{
		w:                w,
		dropAfterPercent: dropAfterPercent,
		dropAfterBytes:   dropAfterBytes,
		expectedSize:     expectedSize,
	}

	// Calculate drop point from percentage if bytes not specified
	if dropAfterBytes <= 0 && dropAfterPercent > 0 && expectedSize > 0 {
		cdw.dropAfterBytes = int64(float64(expectedSize) * dropAfterPercent)
	}

	// Default: drop after 50% if nothing specified
	if cdw.dropAfterBytes <= 0 && expectedSize > 0 {
		cdw.dropAfterBytes = expectedSize / 2
	}

	return cdw
}

// Header returns the header map
func (cdw *ConnectionDropWriter) Header() http.Header {
	return cdw.w.Header()
}

// WriteHeader sends the HTTP response header
func (cdw *ConnectionDropWriter) WriteHeader(statusCode int) {
	if cdw.headersSent.CompareAndSwap(false, true) {
		cdw.w.WriteHeader(statusCode)
	}
}

// Write writes data and may drop the connection mid-write
func (cdw *ConnectionDropWriter) Write(p []byte) (int, error) {
	if cdw.dropped.Load() {
		return 0, errors.New("chaos: connection dropped")
	}

	// Ensure headers are sent
	if !cdw.headersSent.Load() {
		cdw.WriteHeader(http.StatusOK)
	}

	// Check if we should drop before writing
	remaining := cdw.dropAfterBytes - cdw.bytesWritten
	if remaining <= 0 && cdw.dropAfterBytes > 0 {
		cdw.drop()
		return 0, errors.New("chaos: connection dropped")
	}

	// Check if we need to write partial data before dropping
	if cdw.dropAfterBytes > 0 && int64(len(p)) > remaining {
		// Write partial data
		n, _ := cdw.w.Write(p[:remaining])
		cdw.bytesWritten += int64(n)

		// Flush before dropping
		if flusher, ok := cdw.w.(http.Flusher); ok {
			flusher.Flush()
		}

		// Drop connection
		cdw.drop()
		return n, errors.New("chaos: connection dropped")
	}

	// Normal write
	n, err := cdw.w.Write(p)
	cdw.bytesWritten += int64(n)
	return n, err
}

// drop closes the underlying connection
func (cdw *ConnectionDropWriter) drop() {
	if !cdw.dropped.CompareAndSwap(false, true) {
		return // Already dropped
	}

	// Use Hijacker to get raw connection and close it
	if hijacker, ok := cdw.w.(http.Hijacker); ok {
		conn, _, err := hijacker.Hijack()
		if err == nil && conn != nil {
			conn.Close()
		}
	}
}

// Flush implements http.Flusher
func (cdw *ConnectionDropWriter) Flush() {
	if !cdw.dropped.Load() {
		if flusher, ok := cdw.w.(http.Flusher); ok {
			flusher.Flush()
		}
	}
}

// Hijack implements http.Hijacker
func (cdw *ConnectionDropWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := cdw.w.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, errors.New("underlying ResponseWriter does not support hijacking")
}

// IsDropped returns whether the connection has been dropped
func (cdw *ConnectionDropWriter) IsDropped() bool {
	return cdw.dropped.Load()
}

// BytesWritten returns total bytes written
func (cdw *ConnectionDropWriter) BytesWritten() int64 {
	return cdw.bytesWritten
}

// TruncationWriter wraps http.ResponseWriter to truncate response body.
// This simulates partial responses, incomplete downloads, or data loss.
type TruncationWriter struct {
	w               http.ResponseWriter
	truncatePercent float64 // Keep this percentage of body
	maxBytes        int64   // Maximum bytes to write (calculated from percent)
	bytesWritten    int64
	truncated       atomic.Bool
	headersSent     atomic.Bool
}

// NewTruncationWriter creates a new truncation writer
func NewTruncationWriter(w http.ResponseWriter, truncatePercent float64, expectedSize int64) *TruncationWriter {
	tw := &TruncationWriter{
		w:               w,
		truncatePercent: truncatePercent,
	}

	// Calculate max bytes from percentage
	if truncatePercent > 0 && truncatePercent < 1.0 && expectedSize > 0 {
		tw.maxBytes = int64(float64(expectedSize) * truncatePercent)
	}

	return tw
}

// Header returns the header map
func (tw *TruncationWriter) Header() http.Header {
	return tw.w.Header()
}

// WriteHeader sends the HTTP response header
func (tw *TruncationWriter) WriteHeader(statusCode int) {
	if tw.headersSent.CompareAndSwap(false, true) {
		tw.w.WriteHeader(statusCode)
	}
}

// Write writes data up to the truncation limit
func (tw *TruncationWriter) Write(p []byte) (int, error) {
	if tw.truncated.Load() {
		// Silently discard data after truncation
		return len(p), nil
	}

	// Ensure headers are sent
	if !tw.headersSent.Load() {
		tw.WriteHeader(http.StatusOK)
	}

	// Check if we've hit the limit
	if tw.maxBytes > 0 {
		remaining := tw.maxBytes - tw.bytesWritten
		if remaining <= 0 {
			tw.truncated.Store(true)
			return len(p), nil // Pretend we wrote it all
		}

		if int64(len(p)) > remaining {
			// Truncate this write
			p = p[:remaining]
			tw.truncated.Store(true)
		}
	}

	n, err := tw.w.Write(p)
	tw.bytesWritten += int64(n)
	return n, err
}

// Flush implements http.Flusher
func (tw *TruncationWriter) Flush() {
	if flusher, ok := tw.w.(http.Flusher); ok {
		flusher.Flush()
	}
}

// Hijack implements http.Hijacker
func (tw *TruncationWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := tw.w.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, errors.New("underlying ResponseWriter does not support hijacking")
}

// IsTruncated returns whether the response was truncated
func (tw *TruncationWriter) IsTruncated() bool {
	return tw.truncated.Load()
}

// BytesWritten returns total bytes written
func (tw *TruncationWriter) BytesWritten() int64 {
	return tw.bytesWritten
}

// MaxBytes returns the maximum bytes limit
func (tw *TruncationWriter) MaxBytes() int64 {
	return tw.maxBytes
}

// ErrorInjectionWriter wraps http.ResponseWriter to inject HTTP errors.
// Instead of proxying the actual response, it returns an error response.
type ErrorInjectionWriter struct {
	w           http.ResponseWriter
	statusCode  int
	message     string
	injected    atomic.Bool
	headersSent atomic.Bool
}

// NewErrorInjectionWriter creates a new error injection writer
func NewErrorInjectionWriter(w http.ResponseWriter, statusCode int, message string) *ErrorInjectionWriter {
	if message == "" {
		message = http.StatusText(statusCode)
	}
	return &ErrorInjectionWriter{
		w:          w,
		statusCode: statusCode,
		message:    message,
	}
}

// Header returns the header map
func (ew *ErrorInjectionWriter) Header() http.Header {
	return ew.w.Header()
}

// WriteHeader ignores the status code and writes the error status
func (ew *ErrorInjectionWriter) WriteHeader(statusCode int) {
	if ew.headersSent.CompareAndSwap(false, true) {
		// Override with our error status
		ew.w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		ew.w.Header().Set("X-Chaos-Injected", "true")
		ew.w.WriteHeader(ew.statusCode)
	}
}

// Write discards the actual response and writes the error message
func (ew *ErrorInjectionWriter) Write(p []byte) (int, error) {
	if ew.injected.CompareAndSwap(false, true) {
		// First write - send our error message instead
		if !ew.headersSent.Load() {
			ew.WriteHeader(ew.statusCode)
		}
		ew.w.Write([]byte(ew.message))
	}
	// Pretend we wrote it all
	return len(p), nil
}

// Flush implements http.Flusher
func (ew *ErrorInjectionWriter) Flush() {
	if flusher, ok := ew.w.(http.Flusher); ok {
		flusher.Flush()
	}
}

// Hijack implements http.Hijacker
func (ew *ErrorInjectionWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := ew.w.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, errors.New("underlying ResponseWriter does not support hijacking")
}

// StatusCode returns the injected status code
func (ew *ErrorInjectionWriter) StatusCode() int {
	return ew.statusCode
}
