package loadsim

import (
	"context"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

type Runner struct {
	Dial DialFunc
}

type Stats struct {
	ConnectionsOpened int64
	ConnectionsFailed int64
	FramesWritten     int64
	WriteErrors       int64
}

func (r Runner) Run(ctx context.Context, cfg Config) (Stats, error) {
	if _, err := (&FlagConfig{
		protocol:     string(cfg.Protocol),
		addr:         cfg.Addr,
		connections:  cfg.Connections,
		connectRate:  cfg.ConnectRatePerSecond,
		sendInterval: cfg.SendInterval,
		duration:     cfg.Duration,
		template:     cfg.Template,
	}).Build(); err != nil {
		return Stats{}, err
	}
	var factory *FrameFactory
	if cfg.SendFrames {
		var err error
		factory, err = NewFrameFactory(cfg.Protocol, cfg.Template)
		if err != nil {
			return Stats{}, err
		}
	}
	dial := r.Dial
	if dial == nil {
		dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
		dial = dialer.DialContext
	}

	runCtx, cancel := context.WithTimeout(ctx, cfg.Duration)
	defer cancel()

	var stats Stats
	var wg sync.WaitGroup
	connectionIndex := 0
	for i, batch := range ConnectionBatches(cfg.Connections, cfg.ConnectRatePerSecond) {
		if i > 0 {
			if !sleepOrDone(runCtx, time.Second) {
				break
			}
		}
		for range batch {
			if runCtx.Err() != nil {
				break
			}
			conn, err := dial(runCtx, "tcp", cfg.Addr)
			if err != nil {
				atomic.AddInt64(&stats.ConnectionsFailed, 1)
				continue
			}
			atomic.AddInt64(&stats.ConnectionsOpened, 1)
			wg.Add(1)
			connIndex := connectionIndex
			connectionIndex++
			go func() {
				defer wg.Done()
				defer conn.Close()
				if cfg.SendFrames {
					writeLoop(runCtx, conn, factory, connIndex, cfg.SendInterval, &stats)
					return
				}
				<-runCtx.Done()
			}()
		}
	}
	<-runCtx.Done()
	wg.Wait()
	return stats, nil
}

func writeLoop(ctx context.Context, conn net.Conn, factory *FrameFactory, connectionIndex int, interval time.Duration, stats *Stats) {
	for frameIndex := int64(0); ; frameIndex++ {
		payload, err := factory.Frame(connectionIndex, frameIndex)
		if err != nil {
			atomic.AddInt64(&stats.WriteErrors, 1)
			return
		}
		_ = conn.SetWriteDeadline(time.Now().Add(3 * time.Second))
		if _, err := conn.Write(payload); err != nil {
			atomic.AddInt64(&stats.WriteErrors, 1)
			return
		}
		atomic.AddInt64(&stats.FramesWritten, 1)
		if !sleepOrDone(ctx, interval) {
			return
		}
	}
}

func sleepOrDone(ctx context.Context, duration time.Duration) bool {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}