package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"
)

var (
	port     = flag.String("p", "26666", "Port to listen on")
	tcpLog   = flag.String("tcp-log", "tcp.log", "TCP log file path")
	udpLog   = flag.String("udp-log", "udp.log", "UDP log file path")
	httpLog  = flag.String("http-log", "http.log", "HTTP log file path")
	daemon   = flag.Bool("d", false, "Run as daemon")
)

// Logger for each service
var (
	tcpLogger  *log.Logger
	udpLogger  *log.Logger
	httpLogger *log.Logger
)

func main() {
	flag.Parse()

	// If daemon flag is set, restart the program in background
	if *daemon {
		// Get the path of the current executable
		executable, err := os.Executable()
		if err != nil {
			log.Fatal("Failed to get executable path:", err)
		}

		// Prepare command to run the program in background
		// We exclude the -d flag to prevent infinite loop
		args := []string{}
		if *port != "26666" {
			args = append(args, "-p", *port)
		}
		if *tcpLog != "tcp.log" {
			args = append(args, "-tcp-log", *tcpLog)
		}
		if *udpLog != "udp.log" {
			args = append(args, "-udp-log", *udpLog)
		}
		if *httpLog != "http.log" {
			args = append(args, "-http-log", *httpLog)
		}

		cmd := exec.Command(executable, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		// Start the command in background
		err = cmd.Start()
		if err != nil {
			log.Fatal("Failed to start daemon:", err)
		}

		// Print the PID of the background process
		fmt.Printf("Server started as daemon with PID %d\n", cmd.Process.Pid)
		return
	}

	// Open log files
	tcpFile, err := os.OpenFile(*tcpLog, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatal("Failed to open TCP log file:", err)
	}
	defer tcpFile.Close()

	udpFile, err := os.OpenFile(*udpLog, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatal("Failed to open UDP log file:", err)
	}
	defer udpFile.Close()

	httpFile, err := os.OpenFile(*httpLog, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatal("Failed to open HTTP log file:", err)
	}
	defer httpFile.Close()

	// Create loggers for each service
	tcpLogger = log.New(tcpFile, "", 0)
	udpLogger = log.New(udpFile, "", 0)
	httpLogger = log.New(httpFile, "", 0)

	// Start TCP server
	go startTCPServer()

	// Start UDP server
	go startUDPServer()

	// Start HTTP server
	go startHTTPServer()

	// Wait forever
	select {}
}

func startTCPServer() {
	// Start listening on specified port for TCP
	listener, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		log.Fatal("Failed to start TCP server:", err)
	}
	defer listener.Close()

	fmt.Printf("TCP Server listening on port %s\n", *port)
	tcpLogger.Printf("[%s] TCP Server listening on port %s", time.Now().Format("2006-01-02 15:04:05"), *port)

	for {
		// Accept incoming connections
		conn, err := listener.Accept()
		if err != nil {
			tcpLogger.Printf("[%s] Failed to accept connection: %v", time.Now().Format("2006-01-02 15:04:05"), err)
			continue
		}

		// Handle each connection in a separate goroutine
		go handleTCPConnection(conn)
	}
}

func handleTCPConnection(conn net.Conn) {
	defer conn.Close()

	// Get connection info
	remoteAddr := conn.RemoteAddr().(*net.TCPAddr)
	
	// Log when client connects with consistent format
	connectTimestamp := time.Now().Format("2006-01-02 15:04:05")
	connectMsg := fmt.Sprintf("[%s] [TCP] New connection from IP: %s, Port: %d", 
		connectTimestamp, remoteAddr.IP, remoteAddr.Port)
	fmt.Println(connectMsg)
	tcpLogger.Println(connectMsg)

	// Create a scanner to read messages from the client
	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		message := scanner.Text()
		msgTimestamp := time.Now().Format("2006-01-02 15:04:05")

		// Print and log the received message with timestamp
		logEntry := fmt.Sprintf("[%s] [TCP] Received: %s", msgTimestamp, message)
		fmt.Println(logEntry)
		tcpLogger.Println(logEntry)
	}

	if err := scanner.Err(); err != nil {
		errorTimestamp := time.Now().Format("2006-01-02 15:04:05")
		errorMsg := fmt.Sprintf("[%s] [TCP] Error reading from connection: %v", errorTimestamp, err)
		fmt.Println(errorMsg)
		tcpLogger.Println(errorMsg)
	}

	// Log when client disconnects with consistent format
	disconnectTimestamp := time.Now().Format("2006-01-02 15:04:05")
	disconnectMsg := fmt.Sprintf("[%s] [TCP] Client %s:%d disconnected", 
		disconnectTimestamp, remoteAddr.IP, remoteAddr.Port)
	fmt.Println(disconnectMsg)
	tcpLogger.Println(disconnectMsg)
}

func startUDPServer() {
	// Start listening on specified port for UDP
	addr, err := net.ResolveUDPAddr("udp", ":"+*port)
	if err != nil {
		log.Fatal("Failed to resolve UDP address:", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatal("Failed to start UDP server:", err)
	}
	defer conn.Close()

	fmt.Printf("UDP Server listening on port %s\n", *port)
	udpLogger.Printf("[%s] UDP Server listening on port %s", time.Now().Format("2006-01-02 15:04:05"), *port)

	// Buffer for reading data
	buf := make([]byte, 1024)

	for {
		// Read from UDP connection
		n, remoteAddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			udpLogger.Printf("[%s] Failed to read UDP message: %v", time.Now().Format("2006-01-02 15:04:05"), err)
			continue
		}

		// Handle the UDP message
		go handleUDPMessage(remoteAddr, buf[:n])
	}
}

func handleUDPMessage(remoteAddr *net.UDPAddr, data []byte) {
	// Log when client sends data with consistent format
	timestamp := time.Now().Format("2006-01-02 15:04:05")
	connectMsg := fmt.Sprintf("[%s] [UDP] New connection from IP: %s, Port: %d", 
		timestamp, remoteAddr.IP, remoteAddr.Port)
	fmt.Println(connectMsg)
	udpLogger.Println(connectMsg)

	// Log the received message
	message := string(data)
	msgTimestamp := time.Now().Format("2006-01-02 15:04:05")
	logEntry := fmt.Sprintf("[%s] [UDP] Received: %s", msgTimestamp, message)
	fmt.Println(logEntry)
	udpLogger.Println(logEntry)
}

func startHTTPServer() {
	// HTTP server will run on port+1
	httpPort := fmt.Sprintf("%d", parseInt(*port)+1)
	
	// Register ping handler for both GET and POST
	http.HandleFunc("/ping", pingHandler)

	fmt.Printf("HTTP Server listening on port %s\n", httpPort)
	httpLogger.Printf("[%s] HTTP Server listening on port %s", time.Now().Format("2006-01-02 15:04:05"), httpPort)

	// Start HTTP server
	log.Fatal(http.ListenAndServe(":"+httpPort, nil))
}

func pingHandler(w http.ResponseWriter, r *http.Request) {
	// Log the request
	timestamp := time.Now().Format("2006-01-02 15:04:05")
	method := r.Method
	remoteAddr := r.RemoteAddr
	
	// Parse remote address
	var ip string
	var port string
	if strings.Contains(remoteAddr, ":") {
		parts := strings.Split(remoteAddr, ":")
		ip = strings.Join(parts[:len(parts)-1], ":")
		port = parts[len(parts)-1]
	} else {
		ip = remoteAddr
		port = "unknown"
	}
	
	// Log connection
	connectMsg := fmt.Sprintf("[%s] [HTTP] New %s request from IP: %s, Port: %s", 
		timestamp, method, ip, port)
	fmt.Println(connectMsg)
	httpLogger.Println(connectMsg)

	// Read request body for POST requests
	var body string
	if r.Method == "POST" {
		bodyBytes, err := io.ReadAll(r.Body)
		if err != nil {
			errorMsg := fmt.Sprintf("[%s] [HTTP] Error reading request body: %v", 
				time.Now().Format("2006-01-02 15:04:05"), err)
			fmt.Println(errorMsg)
			httpLogger.Println(errorMsg)
		} else {
			body = string(bodyBytes)
		}
	} else if r.Method == "GET" {
		// For GET requests, collect query parameters
		query := r.URL.Query()
		if len(query) > 0 {
			body = fmt.Sprintf("Query params: %v", query)
		}
	}

	// Log the received data
	if body != "" {
		msgTimestamp := time.Now().Format("2006-01-02 15:04:05")
		logEntry := fmt.Sprintf("[%s] [HTTP] Received: %s", msgTimestamp, body)
		fmt.Println(logEntry)
		httpLogger.Println(logEntry)
	}

	// Send response
	response := "pong"
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(response))
	
	// Log response
	responseMsg := fmt.Sprintf("[%s] [HTTP] Sent: %s", time.Now().Format("2006-01-02 15:04:05"), response)
	fmt.Println(responseMsg)
	httpLogger.Println(responseMsg)
}

// Helper function to convert string to int
func parseInt(s string) int {
	var result int
	fmt.Sscanf(s, "%d", &result)
	return result
}