package firewall

import (
	"bufio"
	"bytes"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

// ListActiveConnections executes `conntrack -L -o extended` and parses the output
func ListActiveConnections() ([]ActiveConnection, error) {
	cmd := exec.Command("conntrack", "-L", "-o", "extended")
	out, err := cmd.Output()
	if err != nil {
		// If conntrack is empty or fails, return empty slice
		return make([]ActiveConnection, 0), nil
	}

	var conns []ActiveConnection
	scanner := bufio.NewScanner(bytes.NewReader(out))

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "conntrack") {
			continue
		}

		conn := parseConntrackLine(line)
		if conn.Protocol != "" && conn.SrcAddress != "" {
			conns = append(conns, conn)
		}
	}

	if conns == nil {
		conns = make([]ActiveConnection, 0)
	}
	return conns, nil
}

func parseConntrackLine(line string) ActiveConnection {
	var conn ActiveConnection
	fields := strings.Fields(line)
	if len(fields) < 4 {
		return conn
	}

	// Detect protocol
	protocol := ""
	timeout := 0
	state := ""

	var srcList []string
	var dstList []string
	var sportList []int
	var dportList []int

	for i, f := range fields {
		if f == "tcp" || f == "udp" || f == "icmp" || f == "gre" || f == "esp" {
			protocol = strings.ToUpper(f)
		} else if strings.HasPrefix(f, "src=") {
			srcList = append(srcList, strings.TrimPrefix(f, "src="))
		} else if strings.HasPrefix(f, "dst=") {
			dstList = append(dstList, strings.TrimPrefix(f, "dst="))
		} else if strings.HasPrefix(f, "sport=") {
			p, _ := strconv.Atoi(strings.TrimPrefix(f, "sport="))
			sportList = append(sportList, p)
		} else if strings.HasPrefix(f, "dport=") {
			p, _ := strconv.Atoi(strings.TrimPrefix(f, "dport="))
			dportList = append(dportList, p)
		} else if strings.HasPrefix(f, "mark=") {
			conn.ConnectionMark = strings.TrimPrefix(f, "mark=")
		} else if f == "[ASSURED]" {
			conn.Assured = true
		} else if i > 1 && timeout == 0 {
			if t, err := strconv.Atoi(f); err == nil && t > 0 {
				timeout = t
			}
		}
		if f == "ESTABLISHED" || f == "TIME_WAIT" || f == "CLOSE_WAIT" || f == "FIN_WAIT" || f == "SYN_SENT" || f == "SYN_RECV" {
			state = f
		}
	}

	conn.Protocol = protocol
	conn.Timeout = timeout
	conn.TCPState = state
	if conn.TCPState == "" {
		conn.TCPState = protocol
	}

	if len(srcList) > 0 {
		conn.SrcAddress = srcList[0]
	}
	if len(dstList) > 0 {
		conn.DstAddress = dstList[0]
	}
	if len(sportList) > 0 {
		conn.SrcPort = sportList[0]
	}
	if len(dportList) > 0 {
		conn.DstPort = dportList[0]
	}

	// Reply pair
	if len(srcList) > 1 {
		conn.ReplySrcAddress = srcList[1]
	}
	if len(dstList) > 1 {
		conn.ReplyDstAddress = dstList[1]
	}
	if len(sportList) > 1 {
		conn.ReplySrcPort = sportList[1]
	}
	if len(dportList) > 1 {
		conn.ReplyDstPort = dportList[1]
	}

	conn.ID = fmt.Sprintf("%s-%s:%d-%s:%d", conn.Protocol, conn.SrcAddress, conn.SrcPort, conn.DstAddress, conn.DstPort)
	return conn
}

// KillConnection deletes a specific conntrack flow
func KillConnection(protocol, src string, sport int, dst string, dport int) error {
	args := []string{"-D"}
	if protocol != "" {
		args = append(args, "-p", strings.ToLower(protocol))
	}
	if src != "" {
		args = append(args, "-s", src)
	}
	if dst != "" {
		args = append(args, "-d", dst)
	}
	if sport > 0 {
		args = append(args, "--sport", strconv.Itoa(sport))
	}
	if dport > 0 {
		args = append(args, "--dport", strconv.Itoa(dport))
	}

	cmd := exec.Command("conntrack", args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to terminate connection: %s: %w", string(out), err)
	}
	return nil
}