package main
package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net"
	"os"
	"sync"
	"time"
)

// StressConfig defines the configuration for stress testing
type StressConfig struct {
	Server struct {
		Host string `json:"host"`
		Port int    `json:"port"`
	} `json:"server"`
	
	Test struct {
		MaxClients      int `json:"max_clients"`
		ConnectDelay    int `json:"connect_delay_ms"`
		ActionInterval  int `json:"action_interval_ms"`
		TestDuration    int `json:"test_duration_seconds"`
		RandomSeed      int `json:"random_seed"`
	} `json:"test"`
	
	Behavior struct {
		JoinChannels    bool `json:"join_channels"`
		SendMessages    bool `json:"send_messages"`
		ChangeNicks     bool `json:"change_nicks"`
		UseNewCommands  bool `json:"use_new_commands"`
		RandomQuit      bool `json:"random_quit"`
		
		MessageRate     float64 `json:"message_rate"`
		ChannelJoinRate float64 `json:"channel_join_rate"`
		CommandRate     float64 `json:"command_rate"`
	} `json:"behavior"`
	
	Channels []string `json:"channels"`
	Messages []string `json:"messages"`
	Commands []string `json:"commands"`
}

// IRCClient represents a single IRC client connection
type IRCClient struct {
	ID       int
	Nick     string
	Conn     net.Conn
	Reader   *bufio.Reader
	Writer   *bufio.Writer
	Channels []string
	Active   bool
	mu       sync.Mutex
}

// StressTest manages the entire stress testing operation
type StressTest struct {
	Config  *StressConfig
	Clients []*IRCClient
	Stats   *TestStats
	quit    chan bool
}

// TestStats tracks testing statistics
type TestStats struct {
	ConnectedClients    int
	MessagesSent        int
	CommandsSent        int
	ChannelsJoined      int
	Errors              int
	StartTime           time.Time
	mu                  sync.Mutex
}

// LoadConfig loads configuration from JSON file
func LoadConfig(filename string) (*StressConfig, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("error opening config file: %v", err)
	}
	defer file.Close()

	config := &StressConfig{}
	decoder := json.NewDecoder(file)
	if err := decoder.Decode(config); err != nil {
		return nil, fmt.Errorf("error parsing config: %v", err)
	}

	return config, nil
}

// CreateDefaultConfig creates a default configuration file
func CreateDefaultConfig(filename string) error {
	config := &StressConfig{
		Server: struct {
			Host string `json:"host"`
			Port int    `json:"port"`
		}{
			Host: "localhost",
			Port: 6667,
		},
		Test: struct {
			MaxClients      int `json:"max_clients"`
			ConnectDelay    int `json:"connect_delay_ms"`
			ActionInterval  int `json:"action_interval_ms"`
			TestDuration    int `json:"test_duration_seconds"`
			RandomSeed      int `json:"random_seed"`
		}{
			MaxClients:     300,
			ConnectDelay:   100,
			ActionInterval: 2000,
			TestDuration:   180,
			RandomSeed:     42,
		},
		Behavior: struct {
			JoinChannels    bool `json:"join_channels"`
			SendMessages    bool `json:"send_messages"`
			ChangeNicks     bool `json:"change_nicks"`
			UseNewCommands  bool `json:"use_new_commands"`
			RandomQuit      bool `json:"random_quit"`
			MessageRate     float64 `json:"message_rate"`
			ChannelJoinRate float64 `json:"channel_join_rate"`
			CommandRate     float64 `json:"command_rate"`
		}{
			JoinChannels:    true,
			SendMessages:    true,
			ChangeNicks:     true,
			UseNewCommands:  true,
			RandomQuit:      false,
			MessageRate:     0.4,
			ChannelJoinRate: 0.3,
			CommandRate:     0.2,
		},
		Channels: []string{
			"#test", "#stress", "#general", "#random", "#chaos",
			"#lobby", "#gaming", "#tech", "#chat", "#help",
		},
		Messages: []string{
			"Hello everyone!",
			"This is a stress test message",
			"How is everyone doing?",
			"Testing the server stability",
			"Random message from client",
			"IRC is awesome!",
			"TechIRCd rocks!",
			"Can you see this message?",
			"Stress testing in progress",
			"Everything working fine here",
			"Testing new IRC commands",
			"Server performance looks good",
		},
		Commands: []string{
			"MOTD", "RULES", "MAP", "TIME", "VERSION", "LUSERS",
			"WHO #test", "WHOIS testuser", "LIST", "WHOWAS olduser",
		},
	}

	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("error creating config file: %v", err)
	}
	defer file.Close()

	encoder := json.NewEncoder(file)
	encoder.SetIndent("", "  ")
	return encoder.Encode(config)
}

// NewStressTest creates a new stress test instance
func NewStressTest(config *StressConfig) *StressTest {
	return &StressTest{
		Config:  config,
		Clients: make([]*IRCClient, 0, config.Test.MaxClients),
		Stats: &TestStats{
			StartTime: time.Now(),
		},
		quit: make(chan bool),
	}
}

// Connect establishes connection to IRC server
func (c *IRCClient) Connect(host string, port int) error {
	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return fmt.Errorf("failed to connect: %v", err)
	}

	c.Conn = conn
	c.Reader = bufio.NewReader(conn)
	c.Writer = bufio.NewWriter(conn)
	c.Active = true

	return nil
}

// Register performs IRC client registration
func (c *IRCClient) Register() error {
	commands := []string{
		fmt.Sprintf("NICK %s", c.Nick),
		fmt.Sprintf("USER %s 0 * :Stress Test Client %d", c.Nick, c.ID),
	}

	for _, cmd := range commands {
		if err := c.SendCommand(cmd); err != nil {
			return fmt.Errorf("registration failed: %v", err)
		}
	}

	return nil
}

// SendCommand sends a command to the IRC server
func (c *IRCClient) SendCommand(command string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.Active || c.Writer == nil {
		return fmt.Errorf("client not active")
	}

	_, err := c.Writer.WriteString(command + "\r\n")
	if err != nil {
		return err
	}

	return c.Writer.Flush()
}

// ReadMessages continuously reads messages from server
func (c *IRCClient) ReadMessages(stats *TestStats) {
	defer func() {
		c.Active = false
		if c.Conn != nil {
			c.Conn.Close()
		}
	}()

	for c.Active {
		if c.Reader == nil {
			break
		}

		line, err := c.Reader.ReadString('\n')
		if err != nil {
			stats.mu.Lock()
			stats.Errors++
			stats.mu.Unlock()
			break
		}

		// Handle PING responses
		if len(line) > 4 && line[:4] == "PING" {
			pong := "PONG" + line[4:]
			c.SendCommand(pong[:len(pong)-2]) // Remove \r\n
		}
	}
}

// PerformRandomAction performs a random IRC action
func (c *IRCClient) PerformRandomAction(config *StressConfig, stats *TestStats) {
	if !c.Active {
		return
	}

	action := rand.Float64()

	switch {
	case action < config.Behavior.MessageRate && config.Behavior.SendMessages:
		c.SendRandomMessage(config, stats)
	case action < config.Behavior.MessageRate+config.Behavior.ChannelJoinRate && config.Behavior.JoinChannels:
		c.JoinRandomChannel(config, stats)
	case action < config.Behavior.MessageRate+config.Behavior.ChannelJoinRate+config.Behavior.CommandRate && config.Behavior.UseNewCommands:
		c.SendRandomCommand(config, stats)
	case config.Behavior.ChangeNicks && rand.Float64() < 0.05:
		c.ChangeNick(stats)
	}
}

// SendRandomMessage sends a random message to a random channel
func (c *IRCClient) SendRandomMessage(config *StressConfig, stats *TestStats) {
	if len(c.Channels) == 0 {
		return
	}

	channel := c.Channels[rand.Intn(len(c.Channels))]
	message := config.Messages[rand.Intn(len(config.Messages))]
	
	cmd := fmt.Sprintf("PRIVMSG %s :%s", channel, message)
	if err := c.SendCommand(cmd); err == nil {
		stats.mu.Lock()
		stats.MessagesSent++
		stats.mu.Unlock()
	} else {
		stats.mu.Lock()
		stats.Errors++
		stats.mu.Unlock()
	}
}

// JoinRandomChannel joins a random channel
func (c *IRCClient) JoinRandomChannel(config *StressConfig, stats *TestStats) {
	channel := config.Channels[rand.Intn(len(config.Channels))]
	
	// Check if already in channel
	for _, ch := range c.Channels {
		if ch == channel {
			return
		}
	}

	cmd := fmt.Sprintf("JOIN %s", channel)
	if err := c.SendCommand(cmd); err == nil {
		c.Channels = append(c.Channels, channel)
		stats.mu.Lock()
		stats.ChannelsJoined++
		stats.mu.Unlock()
	} else {
		stats.mu.Lock()
		stats.Errors++
		stats.mu.Unlock()
	}
}

// SendRandomCommand sends a random IRC command
func (c *IRCClient) SendRandomCommand(config *StressConfig, stats *TestStats) {
	command := config.Commands[rand.Intn(len(config.Commands))]
	
	if err := c.SendCommand(command); err == nil {
		stats.mu.Lock()
		stats.CommandsSent++
		stats.mu.Unlock()
	} else {
		stats.mu.Lock()
		stats.Errors++
		stats.mu.Unlock()
	}
}

// ChangeNick changes the client's nickname
func (c *IRCClient) ChangeNick(stats *TestStats) {
	newNick := fmt.Sprintf("User%d_%d", c.ID, rand.Intn(1000))
	cmd := fmt.Sprintf("NICK %s", newNick)
	
	if err := c.SendCommand(cmd); err == nil {
		c.Nick = newNick
	} else {
		stats.mu.Lock()
		stats.Errors++
		stats.mu.Unlock()
	}
}

// CreateClient creates and connects a new IRC client
func (st *StressTest) CreateClient(id int) error {
	client := &IRCClient{
		ID:       id,
		Nick:     fmt.Sprintf("StressUser%d", id),
		Channels: make([]string, 0),
	}

	// Connect to server
	if err := client.Connect(st.Config.Server.Host, st.Config.Server.Port); err != nil {
		return fmt.Errorf("client %d connection failed: %v", id, err)
	}

	// Register with IRC server
	if err := client.Register(); err != nil {
		client.Conn.Close()
		return fmt.Errorf("client %d registration failed: %v", id, err)
	}

	st.Clients = append(st.Clients, client)
	
	st.Stats.mu.Lock()
	st.Stats.ConnectedClients++
	st.Stats.mu.Unlock()

	// Start message reader goroutine
	go client.ReadMessages(st.Stats)

	return nil
}

// RunStressTest executes the complete stress test
func (st *StressTest) RunStressTest() error {
	log.Printf("🚀 Starting TechIRCd stress test with %d clients", st.Config.Test.MaxClients)
	log.Printf("📡 Target server: %s:%d", st.Config.Server.Host, st.Config.Server.Port)
	log.Printf("⏱️  Test duration: %d seconds", st.Config.Test.TestDuration)
	
	// Set random seed
	rand.Seed(int64(st.Config.Test.RandomSeed))

	// Connect clients gradually
	connectDelay := time.Duration(st.Config.Test.ConnectDelay) * time.Millisecond
	for i := 0; i < st.Config.Test.MaxClients; i++ {
		if err := st.CreateClient(i); err != nil {
			log.Printf("❌ Failed to create client %d: %v", i, err)
			st.Stats.mu.Lock()
			st.Stats.Errors++
			st.Stats.mu.Unlock()
			continue
		}

		if i%25 == 0 {
			log.Printf("🔗 Connected %d/%d clients", i+1, st.Config.Test.MaxClients)
		}

		time.Sleep(connectDelay)
	}

	log.Printf("✅ All clients connected! Starting chaos simulation...")

	// Start activity simulation
	actionInterval := time.Duration(st.Config.Test.ActionInterval) * time.Millisecond
	testDuration := time.Duration(st.Config.Test.TestDuration) * time.Second
	
	actionTicker := time.NewTicker(actionInterval)
	defer actionTicker.Stop()

	statsTicker := time.NewTicker(15 * time.Second)
	defer statsTicker.Stop()

	testTimer := time.NewTimer(testDuration)
	defer testTimer.Stop()

	for {
		select {
		case <-actionTicker.C:
			// Perform random actions for random clients
			numActions := rand.Intn(len(st.Clients)/3) + 1
			for i := 0; i < numActions; i++ {
				if len(st.Clients) > 0 {
					clientIndex := rand.Intn(len(st.Clients))
					go st.Clients[clientIndex].PerformRandomAction(st.Config, st.Stats)
				}
			}

		case <-statsTicker.C:
			st.PrintStats()

		case <-testTimer.C:
			log.Println("⏰ Test duration completed. Shutting down...")
			st.Shutdown()
			return nil

		case <-st.quit:
			log.Println("🛑 Test interrupted. Shutting down...")
			st.Shutdown()
			return nil
		}
	}
}

// PrintStats prints current test statistics
func (st *StressTest) PrintStats() {
	st.Stats.mu.Lock()
	defer st.Stats.mu.Unlock()

	elapsed := time.Since(st.Stats.StartTime)
	
	log.Printf("📊 === STRESS TEST STATS ===")
	log.Printf("⏱️  Runtime: %v", elapsed.Round(time.Second))
	log.Printf("👥 Connected Clients: %d", st.Stats.ConnectedClients)
	log.Printf("💬 Messages Sent: %d", st.Stats.MessagesSent)
	log.Printf("⚡ Commands Sent: %d", st.Stats.CommandsSent)
	log.Printf("🏠 Channels Joined: %d", st.Stats.ChannelsJoined)
	log.Printf("❌ Errors: %d", st.Stats.Errors)
	
	if elapsed.Seconds() > 0 {
		log.Printf("📈 Messages/sec: %.2f", float64(st.Stats.MessagesSent)/elapsed.Seconds())
		log.Printf("📈 Commands/sec: %.2f", float64(st.Stats.CommandsSent)/elapsed.Seconds())
	}
	log.Printf("============================")
}

// Shutdown gracefully shuts down all clients
func (st *StressTest) Shutdown() {
	log.Println("🔌 Shutting down all clients...")
	
	for i, client := range st.Clients {
		if client.Active {
			client.SendCommand("QUIT :Stress test completed")
			client.Active = false
			if client.Conn != nil {
				client.Conn.Close()
			}
		}
		
		if i%50 == 0 && i > 0 {
			log.Printf("🔌 Disconnected %d/%d clients", i+1, len(st.Clients))
		}
	}
	
	st.PrintStats()
	log.Println("🎉 Stress test completed successfully!")
}

func main() {
	configFile := "stress_test_config.json"
	
	log.Println("🎯 TechIRCd Go Stress Tester v1.0")
	log.Println("==================================")
	
	// Check if config file exists, create default if not
	if _, err := os.Stat(configFile); os.IsNotExist(err) {
		log.Printf("📝 Config file %s not found. Creating default...", configFile)
		if err := CreateDefaultConfig(configFile); err != nil {
			log.Fatalf("❌ Failed to create default config: %v", err)
		}
		log.Printf("✅ Default config created at %s", configFile)
		log.Println("📖 Please review and modify the config as needed, then run again.")
		return
	}

	// Load configuration
	config, err := LoadConfig(configFile)
	if err != nil {
		log.Fatalf("❌ Failed to load config: %v", err)
	}

	log.Printf("📋 Loaded config: %d clients, %d second test", config.Test.MaxClients, config.Test.TestDuration)

	// Create and run stress test
	stressTest := NewStressTest(config)
	
	if err := stressTest.RunStressTest(); err != nil {
		log.Fatalf("❌ Stress test failed: %v", err)
	}
}