package main

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"text/tabwriter"
	"time"

	"github.com/miekg/dns"
)

type ListenerInfo struct {
	Name       string
	IP         []string // Hosts
	BindIP     string   // Always 0.0.0.0
	Port       string
	Protocol   string
	Listener   net.Listener
	Server     *dns.Server
	HTTPServer *http.Server // For HTTPS listeners
}

var listeners sync.Map

func StartAgentListener(protocol string, hostSlice []string, port, name string, operatorConn net.Conn) {
	// Always bind to 0.0.0.0
	bindIP := "0.0.0.0"
	IpPort := fmt.Sprintf("%s:%s", bindIP, port)

	// Start the listener
	var listener net.Listener
	var httpServer *http.Server
	var server *dns.Server

	var err error

	switch protocol {
	case "tcp":
		listener, err = net.Listen("tcp", IpPort)
	case "ssl":
		listener, err = tls.Listen("tcp", IpPort, tlsConfig)
	case "dns":
		server, err = StartDNSlistener(port)
	case "https":
		httpServer = StartHttpsListener(IpPort)
	default:
		debugLog("Unsupported protocol: %s", protocol)
		operatorConn.Write([]byte(fmt.Sprintf("Unsupported protocol: %s", protocol)))
		return
	}

	// Check for errors
	if err != nil {
		debugLog("Unable to start %s listener: %v", protocol, err)
		operatorConn.Write([]byte(fmt.Sprintf("Unable to start %s listener: %v", protocol, err)))
		return
	}

	// Log success and notify operator
	message := fmt.Sprintf("Listener '%s' started on %s (%s)", name, IpPort, protocol)
	log.Println(message)
	if operatorConn != nil {
		_, err := operatorConn.Write([]byte(message))
		if err != nil {
			debugLog("Failed to notify operator: %v", err)
		}
	}

	// Start goroutine for TCP/SSL listeners or add to list DNS/HTTP servers
	if protocol == "ssl" {
		go AcceptAgentConnections(listener, name)
		AddListener(name, hostSlice, bindIP, port, protocol, listener, nil, nil)
	} else if protocol == "tcp" {
		go AcceptAgentConnections(listener, name)
		AddListener(name, hostSlice, bindIP, port, protocol, listener, nil, nil)
	} else if protocol == "dns" {
		AddListener(name, hostSlice, bindIP, port, protocol, nil, server, nil)
	} else if protocol == "https" {
		AddListener(name, hostSlice, bindIP, port, protocol, nil, nil, httpServer)
	}
}

// Stop listener and delete from list
func StopListener(name string, operatorConn net.Conn) {
	if value, ok := listeners.Load(name); ok {
		listenerInfo := value.(*ListenerInfo)

		// Stop DNS server if present
		if listenerInfo.Server != nil {
			if err := listenerInfo.Server.Shutdown(); err != nil {
				debugLog("Error stopping DNS listener '%s': %v", name, err)
			}
		}

		// Stop HTTPS server if present
		if listenerInfo.HTTPServer != nil {
			if err := listenerInfo.HTTPServer.Close(); err != nil {
				debugLog("Error stopping HTTPS listener '%s': %v", name, err)
			}
		}

		// Stop TCP/TLS listener if present
		if listenerInfo.Listener != nil {
			if err := listenerInfo.Listener.Close(); err != nil {
				debugLog("Error stopping TCP or TLS listener '%s': %v", name, err)
			}
		}

		// Wait a moment for everything to clean up
		time.Sleep(100 * time.Millisecond)

		listeners.Delete(name)

		message := fmt.Sprintf("Listener '%s' stopped successfully", name)
		log.Println(message)

		if operatorConn != nil {
			operatorConn.Write([]byte(message))
		}

		// Broadcast updated listener list
		UpdateListenersTabCompletion()

		return
	}

	// If something wrong - log error
	message := fmt.Sprintf("Listener '%s' not found", name)
	log.Println(message)
	if operatorConn != nil {
		operatorConn.Write([]byte(message))
	}
}

// Add listener to list
func AddListener(name string, hosts []string, bindIP, port, protocol string, listener net.Listener, DNSserver *dns.Server, httpServer *http.Server) {
	listenerInfo := &ListenerInfo{
		Name:       name,
		IP:         hosts,
		BindIP:     bindIP,
		Port:       port,
		Protocol:   protocol,
		Listener:   listener,
		Server:     DNSserver,
		HTTPServer: httpServer,
	}

	listeners.Store(name, listenerInfo)

	// Broadcast updated listener list to operators
	UpdateListenersTabCompletion()
}

// List all active listeners
func ListListeners(operatorID string) {
	operatorConn, _ := GetOperatorConn(operatorID)

	var sb strings.Builder
	tw := tabwriter.NewWriter(&sb, 0, 8, 2, ' ', 0) // Initialize a tab writer

	// Write the header row
	fmt.Fprintln(tw, "Name\tProtocol\tIP(s)\tBind IP\tPort")
	fmt.Fprintln(tw, "==========\t========\t=========\t=========\t=====")

	// Iterate over the listeners and write their details
	listeners.Range(func(key, value interface{}) bool {
		info := value.(*ListenerInfo)

		// Join the IP slice into a comma-separated string
		ips := strings.Join(info.IP, ", ")

		// Write the listener details to the tab writer
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\n",
			info.Name,
			info.Protocol,
			ips,
			info.BindIP,
			info.Port,
		)
		return true
	})

	// Flush the tab writer and write to operator connection
	tw.Flush()

	if sb.Len() == 0 {
		operatorConn.Write([]byte("No active listeners"))
	} else {
		operatorConn.Write([]byte("Active listeners:\n" + sb.String()))
	}
}

// Get listeners list to send to operators to update their tab completion
func GetListenerNames() []string {
	var names []string
	listeners.Range(func(key, value any) bool {
		names = append(names, key.(string))
		return true
	})
	return names
}
