package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

// Color constants for terminal output
const (
	Red    = "\033[0;31m"
	Green  = "\033[0;32m"
	Yellow = "\033[1;33m"
	Blue   = "\033[0;34m"
	Cyan   = "\033[0;36m"
	Bold   = "\033[1m"
	NC     = "\033[0m" // No Color
)

// ProcessInfo holds information about a process using a port
type ProcessInfo struct {
	PID         int
	ProcessName string
	Protocol    string
	DockerInfo  string
}

// DockerContainer represents a Docker container
type DockerContainer struct {
	Name    string
	Image   string
	Ports   []PortMapping
	Network string
}

// PortMapping represents a port mapping
type PortMapping struct {
	ContainerPort int
	HostPort      int
	Protocol      string
	IPv6          bool
}

func main() {
	if len(os.Args) < 2 {
		showUsage()
		os.Exit(1)
	}

	arg := os.Args[1]

	switch arg {
	case "--help", "-h":
		showHelp()
	case "--list", "-l":
		listDockerServices()
	default:
		port, err := strconv.Atoi(arg)
		if err != nil || port < 1 || port > 65535 {
			fmt.Printf("%sError:%s Invalid port number. Must be between 1 and 65535.\n", Red, NC)
			os.Exit(1)
		}
		checkPort(port)
	}
}

func showUsage() {
	fmt.Printf("%sUsage:%s inuse <port_number>\n", Red, NC)
	fmt.Printf("%s       inuse --list%s\n", Yellow, NC)
	fmt.Printf("%s       inuse --help%s\n", Yellow, NC)
	fmt.Printf("%sExample:%s inuse 80\n", Yellow, NC)
	fmt.Printf("%s         inuse --list%s\n", Yellow, NC)
}

func showHelp() {
	fmt.Printf("%s%sinuse - Check if a port is in use%s\n\n", Cyan, Bold, NC)
	fmt.Printf("%sUSAGE:%s\n", Bold, NC)
	fmt.Printf("  inuse <port_number>    Check if a specific port is in use\n")
	fmt.Printf("  inuse --list, -l       List all Docker services with listening ports\n")
	fmt.Printf("  inuse --help, -h       Show this help message\n\n")
	fmt.Printf("%sEXAMPLES:%s\n", Bold, NC)
	fmt.Printf("  %sinuse 80%s              Check if port 80 is in use\n", Green, NC)
	fmt.Printf("  %sinuse 3000%s            Check if port 3000 is in use\n", Green, NC)
	fmt.Printf("  %sinuse --list%s          Show all Docker services with ports\n\n", Green, NC)
	fmt.Printf("%sDESCRIPTION:%s\n", Bold, NC)
	fmt.Printf("  The inuse function checks if a specific port is in use and identifies\n")
	fmt.Printf("  the process using it. It can detect regular processes, Docker containers\n")
	fmt.Printf("  with published ports, and containers using host networking.\n\n")
	fmt.Printf("%sOUTPUT:%s\n", Bold, NC)
	fmt.Printf("  %s✓%s Port is in use - shows process name, PID, and Docker info if applicable\n", Green, NC)
	fmt.Printf("  %s✗%s Port is free\n", Red, NC)
	fmt.Printf("  %s⚠%s Port is in use but process cannot be identified\n", Yellow, NC)
}

func listDockerServices() {
	if !isDockerAvailable() {
		fmt.Printf("%sError:%s Docker is not available\n", Red, NC)
		os.Exit(1)
	}

	fmt.Printf("%s%sDocker Services with Listening Ports:%s\n\n", Cyan, Bold, NC)

	containers := getRunningContainers()
	if len(containers) == 0 {
		fmt.Printf("%sNo running Docker containers found%s\n", Yellow, NC)
		return
	}

	foundServices := false
	for _, container := range containers {
		if len(container.Ports) > 0 {
			cleanImage := cleanImageName(container.Image)
			fmt.Printf("%s📦 %s%s%s %s(%s)%s\n", Green, Bold, container.Name, NC, Cyan, cleanImage, NC)

			for _, port := range container.Ports {
				ipv6Marker := ""
				if port.IPv6 {
					ipv6Marker = " [IPv6]"
				}
				fmt.Printf("%s  ├─ Port %s%d%s%s → %d (%s)%s%s\n",
					Cyan, Bold, port.HostPort, NC, Cyan, port.ContainerPort, port.Protocol, ipv6Marker, NC)
			}
			fmt.Println()
			foundServices = true
		}
	}

	// Check for host networking containers
	hostContainers := getHostNetworkingContainers()
	if len(hostContainers) > 0 {
		fmt.Printf("%s%sHost Networking Containers:%s\n", Yellow, Bold, NC)
		for _, container := range hostContainers {
			cleanImage := cleanImageName(container.Image)
			fmt.Printf("%s🌐 %s%s%s %s(%s)%s %s- uses host networking%s\n",
				Yellow, Bold, container.Name, NC, Cyan, cleanImage, NC, Yellow, NC)
		}
		fmt.Println()
		foundServices = true
	}

	if !foundServices {
		fmt.Printf("%sNo Docker services with exposed ports found%s\n", Yellow, NC)
	}
}

func checkPort(port int) {
	// Check if port is in use first
	if !isPortInUse(port) {
		fmt.Printf("%s✗ Port %d is FREE%s\n", Red, port, NC)
		os.Exit(1)
	}

	// Port is in use, now find what's using it
	process := findProcessUsingPort(port)
	if process != nil {
		dockerInfo := ""
		if process.DockerInfo != "" {
			dockerInfo = " " + process.DockerInfo
		}
		fmt.Printf("%s✓ Port %d (%s) in use by %s%s%s %sas PID %s%d%s%s\n",
			Green, port, process.Protocol, Bold, process.ProcessName, NC, Green, Bold, process.PID, NC, dockerInfo)
		return
	}

	// Check if it's a Docker container
	containerInfo := findDockerContainerUsingPort(port)
	if containerInfo != "" {
		fmt.Printf("%s✓ Port %d in use by Docker container %s\n", Green, port, containerInfo)
		return
	}

	// If we still haven't found the process, check for host networking containers more thoroughly
	hostNetworkProcess := findHostNetworkingProcess(port)
	if hostNetworkProcess != "" {
		fmt.Printf("%s✓ Port %d likely in use by %s\n", Green, port, hostNetworkProcess)
		return
	}

	// If we still haven't found the process
	fmt.Printf("%s⚠ Port %d is in use but unable to identify the process%s\n", Yellow, port, NC)

	if isDockerAvailable() {
		hostContainers := getHostNetworkingContainers()
		if len(hostContainers) > 0 {
			fmt.Printf("%s  Note: Found Docker containers using host networking:%s\n", Cyan, NC)
			for _, container := range hostContainers {
				cleanImage := cleanImageName(container.Image)
				fmt.Printf("%s    - %s (%s)%s\n", Cyan, container.Name, cleanImage, NC)
			}
			fmt.Printf("%s  These containers share the host's network, so one of them might be using this port%s\n", Cyan, NC)
		} else {
			fmt.Printf("%s  This might be due to insufficient permissions or the process being in a different namespace%s\n", Cyan, NC)
		}
	} else {
		fmt.Printf("%s  This might be due to insufficient permissions or the process being in a different namespace%s\n", Cyan, NC)
	}
}

func isPortInUse(port int) bool {
	// Try ss first
	if isCommandAvailable("ss") {
		cmd := exec.Command("ss", "-tulpn")
		output, err := cmd.Output()
		if err == nil {
			portPattern := fmt.Sprintf(":%d ", port)
			return strings.Contains(string(output), portPattern)
		}
	}

	// Try netstat as fallback
	if isCommandAvailable("netstat") {
		cmd := exec.Command("netstat", "-tulpn")
		output, err := cmd.Output()
		if err == nil {
			portPattern := fmt.Sprintf(":%d ", port)
			return strings.Contains(string(output), portPattern)
		}
	}

	return false
}

func findProcessUsingPort(port int) *ProcessInfo {
	// Method 1: Try netstat
	if process := tryNetstat(port); process != nil {
		return process
	}

	// Method 2: Try ss
	if process := trySS(port); process != nil {
		return process
	}

	// Method 3: Try lsof
	if process := tryLsof(port); process != nil {
		return process
	}

	// Method 4: Try fuser
	if process := tryFuser(port); process != nil {
		return process
	}

	return nil
}

func tryNetstat(port int) *ProcessInfo {
	if !isCommandAvailable("netstat") {
		return nil
	}

	cmd := exec.Command("netstat", "-tulpn")
	output, err := cmd.Output()
	if err != nil {
		// Try with sudo if available
		if isCommandAvailable("sudo") {
			cmd = exec.Command("sudo", "netstat", "-tulpn")
			output, err = cmd.Output()
			if err != nil {
				return nil
			}
		} else {
			return nil
		}
	}

	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	portPattern := fmt.Sprintf(":%d ", port)

	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, portPattern) {
			fields := strings.Fields(line)
			if len(fields) >= 7 {
				pidProcess := fields[6]
				parts := strings.Split(pidProcess, "/")
				if len(parts) >= 2 {
					if pid, err := strconv.Atoi(parts[0]); err == nil {
						processName := parts[1]
						protocol := fields[0]
						dockerInfo := getDockerInfo(pid, processName, port)
						return &ProcessInfo{
							PID:         pid,
							ProcessName: processName,
							Protocol:    protocol,
							DockerInfo:  dockerInfo,
						}
					}
				}
			}
		}
	}

	return nil
}

func trySS(port int) *ProcessInfo {
	if !isCommandAvailable("ss") {
		return nil
	}

	cmd := exec.Command("ss", "-tulpn")
	output, err := cmd.Output()
	if err != nil {
		// Try with sudo if available
		if isCommandAvailable("sudo") {
			cmd = exec.Command("sudo", "ss", "-tulpn")
			output, err = cmd.Output()
			if err != nil {
				return nil
			}
		} else {
			return nil
		}
	}

	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	portPattern := fmt.Sprintf(":%d ", port)
	pidRegex := regexp.MustCompile(`pid=(\d+)`)

	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, portPattern) {
			matches := pidRegex.FindStringSubmatch(line)
			if len(matches) >= 2 {
				if pid, err := strconv.Atoi(matches[1]); err == nil {
					processName := getProcessName(pid)
					if processName != "" {
						fields := strings.Fields(line)
						protocol := ""
						if len(fields) > 0 {
							protocol = fields[0]
						}
						dockerInfo := getDockerInfo(pid, processName, port)
						return &ProcessInfo{
							PID:         pid,
							ProcessName: processName,
							Protocol:    protocol,
							DockerInfo:  dockerInfo,
						}
					}
				}
			}
		}
	}

	return nil
}

func tryLsof(port int) *ProcessInfo {
	if !isCommandAvailable("lsof") {
		return nil
	}

	cmd := exec.Command("lsof", "-i", fmt.Sprintf(":%d", port), "-n", "-P")
	output, err := cmd.Output()
	if err != nil {
		// Try with sudo if available
		if isCommandAvailable("sudo") {
			cmd = exec.Command("sudo", "lsof", "-i", fmt.Sprintf(":%d", port), "-n", "-P")
			output, err = cmd.Output()
			if err != nil {
				return nil
			}
		} else {
			return nil
		}
	}

	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "LISTEN") {
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				processName := fields[0]
				if pid, err := strconv.Atoi(fields[1]); err == nil {
					dockerInfo := getDockerInfo(pid, processName, port)
					return &ProcessInfo{
						PID:         pid,
						ProcessName: processName,
						Protocol:    "tcp",
						DockerInfo:  dockerInfo,
					}
				}
			}
		}
	}

	return nil
}

func tryFuser(port int) *ProcessInfo {
	if !isCommandAvailable("fuser") {
		return nil
	}

	cmd := exec.Command("fuser", fmt.Sprintf("%d/tcp", port))
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	pids := strings.Fields(string(output))
	for _, pidStr := range pids {
		if pid, err := strconv.Atoi(strings.TrimSpace(pidStr)); err == nil {
			processName := getProcessName(pid)
			if processName != "" {
				return &ProcessInfo{
					PID:         pid,
					ProcessName: processName,
					Protocol:    "tcp",
					DockerInfo:  "",
				}
			}
		}
	}

	return nil
}

func getProcessName(pid int) string {
	cmd := exec.Command("ps", "-p", strconv.Itoa(pid), "-o", "comm=")
	output, err := cmd.Output()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(output))
}

func getDockerInfo(pid int, processName string, port int) string {
	if !isDockerAvailable() {
		return ""
	}

	// Check if it's docker-proxy (handle truncated names like "docker-pr")
	if processName == "docker-proxy" || strings.HasPrefix(processName, "docker-pr") {
		containerName := getContainerByPublishedPort(port)
		if containerName != "" {
			image := getContainerImage(containerName)
			cleanImage := cleanImageName(image)
			return fmt.Sprintf("%s(Docker: %s, image: %s)%s", Cyan, containerName, cleanImage, NC)
		}
		return fmt.Sprintf("%s(Docker proxy)%s", Cyan, NC)
	}

	// Check if process is in a Docker container using cgroup
	containerInfo := getContainerByPID(pid)
	if containerInfo != "" {
		return fmt.Sprintf("%s(Docker: %s)%s", Cyan, containerInfo, NC)
	}

	// Check if this process might be in a host networking container
	hostContainer := checkHostNetworkingContainer(pid, processName)
	if hostContainer != "" {
		return fmt.Sprintf("%s(Docker host network: %s)%s", Cyan, hostContainer, NC)
	}

	return ""
}

func getContainerByPID(pid int) string {
	cgroupPath := fmt.Sprintf("/proc/%d/cgroup", pid)
	file, err := os.Open(cgroupPath)
	if err != nil {
		return ""
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	containerIDRegex := regexp.MustCompile(`[a-f0-9]{64}`)

	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "docker") {
			matches := containerIDRegex.FindStringSubmatch(line)
			if len(matches) > 0 {
				containerID := matches[0]
				containerName := getContainerNameByID(containerID)
				if containerName != "" {
					return containerName
				}
				return containerID[:12]
			}
		}
	}

	return ""
}

func findDockerContainerUsingPort(port int) string {
	if !isDockerAvailable() {
		return ""
	}

	// Check for containers with published ports
	cmd := exec.Command("docker", "ps", "--format", "{{.Names}}", "--filter", fmt.Sprintf("publish=%d", port))
	output, err := cmd.Output()
	if err != nil {
		return ""
	}

	containerName := strings.TrimSpace(string(output))
	if containerName != "" {
		image := getContainerImage(containerName)
		cleanImage := cleanImageName(image)
		return fmt.Sprintf("%s%s%s %s(published port, image: %s)%s", Bold, containerName, NC, Cyan, cleanImage, NC)
	}

	return ""
}

func isDockerAvailable() bool {
	return isCommandAvailable("docker")
}

func isCommandAvailable(command string) bool {
	_, err := exec.LookPath(command)
	return err == nil
}

func getRunningContainers() []DockerContainer {
	if !isDockerAvailable() {
		return nil
	}

	cmd := exec.Command("docker", "ps", "--format", "{{.Names}}")
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	var containers []DockerContainer
	scanner := bufio.NewScanner(strings.NewReader(string(output)))

	for scanner.Scan() {
		containerName := strings.TrimSpace(scanner.Text())
		if containerName != "" {
			container := DockerContainer{
				Name:  containerName,
				Image: getContainerImage(containerName),
				Ports: getContainerPorts(containerName),
			}
			containers = append(containers, container)
		}
	}

	return containers
}

func getHostNetworkingContainers() []DockerContainer {
	if !isDockerAvailable() {
		return nil
	}

	cmd := exec.Command("docker", "ps", "--format", "{{.Names}}", "--filter", "network=host")
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	var containers []DockerContainer
	scanner := bufio.NewScanner(strings.NewReader(string(output)))

	for scanner.Scan() {
		containerName := strings.TrimSpace(scanner.Text())
		if containerName != "" {
			container := DockerContainer{
				Name:    containerName,
				Image:   getContainerImage(containerName),
				Network: "host",
			}
			containers = append(containers, container)
		}
	}

	return containers
}

func getContainerImage(containerName string) string {
	cmd := exec.Command("docker", "inspect", containerName)
	output, err := cmd.Output()
	if err != nil {
		return ""
	}

	var inspectData []map[string]interface{}
	if err := json.Unmarshal(output, &inspectData); err != nil {
		return ""
	}

	if len(inspectData) > 0 {
		if image, ok := inspectData[0]["Config"].(map[string]interface{})["Image"].(string); ok {
			return image
		}
	}

	return ""
}

func getContainerPorts(containerName string) []PortMapping {
	cmd := exec.Command("docker", "port", containerName)
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	var ports []PortMapping
	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	portRegex := regexp.MustCompile(`(\d+)/(tcp|udp) -> 0\.0\.0\.0:(\d+)`)
	ipv6PortRegex := regexp.MustCompile(`(\d+)/(tcp|udp) -> \[::\]:(\d+)`)

	for scanner.Scan() {
		line := scanner.Text()

		// Check for IPv4
		if matches := portRegex.FindStringSubmatch(line); len(matches) >= 4 {
			containerPort, _ := strconv.Atoi(matches[1])
			protocol := matches[2]
			hostPort, _ := strconv.Atoi(matches[3])

			ports = append(ports, PortMapping{
				ContainerPort: containerPort,
				HostPort:      hostPort,
				Protocol:      protocol,
				IPv6:          false,
			})
		}

		// Check for IPv6
		if matches := ipv6PortRegex.FindStringSubmatch(line); len(matches) >= 4 {
			containerPort, _ := strconv.Atoi(matches[1])
			protocol := matches[2]
			hostPort, _ := strconv.Atoi(matches[3])

			ports = append(ports, PortMapping{
				ContainerPort: containerPort,
				HostPort:      hostPort,
				Protocol:      protocol,
				IPv6:          true,
			})
		}
	}

	return ports
}

func getContainerByPublishedPort(port int) string {
	cmd := exec.Command("docker", "ps", "--format", "{{.Names}}", "--filter", fmt.Sprintf("publish=%d", port))
	output, err := cmd.Output()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(output))
}

func getContainerNameByID(containerID string) string {
	cmd := exec.Command("docker", "inspect", containerID)
	output, err := cmd.Output()
	if err != nil {
		return ""
	}

	var inspectData []map[string]interface{}
	if err := json.Unmarshal(output, &inspectData); err != nil {
		return ""
	}

	if len(inspectData) > 0 {
		if name, ok := inspectData[0]["Name"].(string); ok {
			return strings.TrimPrefix(name, "/")
		}
	}

	return ""
}

func cleanImageName(image string) string {
	// Remove SHA256 hashes
	shaRegex := regexp.MustCompile(`sha256:[a-f0-9]*`)
	cleaned := shaRegex.ReplaceAllString(image, "[image-hash]")

	// Remove registry prefixes, keep only the last part
	parts := strings.Split(cleaned, "/")
	if len(parts) > 0 {
		return parts[len(parts)-1]
	}

	return cleaned
}

func findHostNetworkingProcess(port int) string {
	if !isDockerAvailable() {
		return ""
	}

	// Get all host networking containers
	hostContainers := getHostNetworkingContainers()

	for _, container := range hostContainers {
		// Check if this container might be using the port
		if isContainerUsingPort(container.Name, port) {
			cleanImage := cleanImageName(container.Image)
			return fmt.Sprintf("%s%s%s %s(Docker host network: %s)%s", Bold, container.Name, NC, Cyan, cleanImage, NC)
		}
	}

	return ""
}

func isContainerUsingPort(containerName string, port int) bool {
	// Try to execute netstat inside the container to see if it's listening on the port
	cmd := exec.Command("docker", "exec", containerName, "sh", "-c",
		fmt.Sprintf("netstat -tlnp 2>/dev/null | grep ':%d ' || ss -tlnp 2>/dev/null | grep ':%d '", port, port))
	output, err := cmd.Output()
	if err != nil {
		return false
	}

	return len(output) > 0
}

func checkHostNetworkingContainer(pid int, processName string) string {
	if !isDockerAvailable() {
		return ""
	}

	// Get all host networking containers and check if any match this process
	hostContainers := getHostNetworkingContainers()

	for _, container := range hostContainers {
		// Try to find this process inside the container
		cmd := exec.Command("docker", "exec", container.Name, "sh", "-c",
			fmt.Sprintf("ps -o pid,comm | grep '%s' | grep -q '%d\\|%s'", processName, pid, processName))
		err := cmd.Run()
		if err == nil {
			cleanImage := cleanImageName(container.Image)
			return fmt.Sprintf("%s (%s)", container.Name, cleanImage)
		}
	}

	return ""
}