package security

import (
	"fmt"
	"os/exec"
	"runtime"
	"strings"
)

// KillSwitch управляет блокировкой интернета при разрыве VPN
type KillSwitch struct {
	enabled       bool
	originalRules []string
}

// NewKillSwitch создает новый Kill Switch
func NewKillSwitch() *KillSwitch {
	return &KillSwitch{
		enabled:       false,
		originalRules: []string{},
	}
}

// Enable включает Kill Switch (блокирует весь трафик кроме VPN)
func (ks *KillSwitch) Enable(vpnInterface string, allowedIPs []string) error {
	if runtime.GOOS != "windows" {
		return ks.enableUnix(vpnInterface, allowedIPs)
	}
	return ks.enableWindows(vpnInterface, allowedIPs)
}

// Disable отключает Kill Switch (восстанавливает нормальный трафик)
func (ks *KillSwitch) Disable() error {
	if runtime.GOOS != "windows" {
		return ks.disableUnix()
	}
	return ks.disableWindows()
}

// enableWindows включает Kill Switch на Windows через Windows Firewall
func (ks *KillSwitch) enableWindows(vpnInterface string, allowedIPs []string) error {
	// Создаем правило блокировки всего исходящего трафика
	cmd := exec.Command("netsh", "advfirewall", "firewall", "add", "rule",
		"name=VPN_KillSwitch_Block_Out",
		"dir=out",
		"action=block",
		"enable=yes")
	
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ошибка создания правила блокировки: %w", err)
	}

	// Создаем правило блокировки всего входящего трафика
	cmd = exec.Command("netsh", "advfirewall", "firewall", "add", "rule",
		"name=VPN_KillSwitch_Block_In",
		"dir=in",
		"action=block",
		"enable=yes")
	
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ошибка создания правила блокировки: %w", err)
	}

	// Разрешаем локальный трафик (127.0.0.1)
	cmd = exec.Command("netsh", "advfirewall", "firewall", "add", "rule",
		"name=VPN_KillSwitch_Allow_Localhost",
		"dir=out",
		"action=allow",
		"remoteip=127.0.0.1",
		"enable=yes")
	
	cmd.Run() // Игнорируем ошибку

	// Разрешаем трафик к VPN серверам
	for _, ip := range allowedIPs {
		cmd = exec.Command("netsh", "advfirewall", "firewall", "add", "rule",
			fmt.Sprintf("name=VPN_KillSwitch_Allow_%s", strings.ReplaceAll(ip, ".", "_")),
			"dir=out",
			"action=allow",
			fmt.Sprintf("remoteip=%s", ip),
			"enable=yes")
		
		cmd.Run() // Игнорируем ошибки для отдельных IP
	}

	ks.enabled = true
	return nil
}

// disableWindows отключает Kill Switch на Windows
func (ks *KillSwitch) disableWindows() error {
	if !ks.enabled {
		return nil
	}

	// Удаляем все правила Kill Switch
	rules := []string{
		"VPN_KillSwitch_Block_Out",
		"VPN_KillSwitch_Block_In",
		"VPN_KillSwitch_Allow_Localhost",
	}

	for _, rule := range rules {
		cmd := exec.Command("netsh", "advfirewall", "firewall", "delete", "rule",
			fmt.Sprintf("name=%s", rule))
		cmd.Run() // Игнорируем ошибки
	}

	// Удаляем правила для разрешенных IP (пытаемся удалить все возможные)
	// Получаем список всех правил с префиксом VPN_KillSwitch_Allow_
	listCmd := exec.Command("netsh", "advfirewall", "firewall", "show", "rule", "name=all")
	output, _ := listCmd.Output()
	
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		if strings.Contains(line, "VPN_KillSwitch_Allow_") {
			parts := strings.Split(line, ":")
			if len(parts) >= 2 {
				ruleName := strings.TrimSpace(parts[1])
				deleteCmd := exec.Command("netsh", "advfirewall", "firewall", "delete", "rule",
					fmt.Sprintf("name=%s", ruleName))
				deleteCmd.Run()
			}
		}
	}

	ks.enabled = false
	return nil
}

// enableUnix включает Kill Switch на Unix системах через iptables
func (ks *KillSwitch) enableUnix(vpnInterface string, allowedIPs []string) error {
	// Сохраняем текущие правила
	cmd := exec.Command("iptables-save")
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("ошибка сохранения правил iptables: %w", err)
	}
	ks.originalRules = strings.Split(string(output), "\n")

	// Блокируем весь исходящий трафик по умолчанию
	cmd = exec.Command("iptables", "-P", "OUTPUT", "DROP")
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ошибка установки политики OUTPUT: %w", err)
	}

	// Разрешаем локальный трафик
	cmd = exec.Command("iptables", "-A", "OUTPUT", "-o", "lo", "-j", "ACCEPT")
	cmd.Run()

	// Разрешаем трафик через VPN интерфейс
	if vpnInterface != "" {
		cmd = exec.Command("iptables", "-A", "OUTPUT", "-o", vpnInterface, "-j", "ACCEPT")
		cmd.Run()
	}

	// Разрешаем трафик к VPN серверам
	for _, ip := range allowedIPs {
		cmd = exec.Command("iptables", "-A", "OUTPUT", "-d", ip, "-j", "ACCEPT")
		cmd.Run()
	}

	ks.enabled = true
	return nil
}

// disableUnix отключает Kill Switch на Unix системах
func (ks *KillSwitch) disableUnix() error {
	if !ks.enabled {
		return nil
	}

	// Восстанавливаем политику по умолчанию
	cmd := exec.Command("iptables", "-P", "OUTPUT", "ACCEPT")
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ошибка восстановления политики: %w", err)
	}

	// Очищаем цепочку OUTPUT
	cmd = exec.Command("iptables", "-F", "OUTPUT")
	cmd.Run()

	ks.enabled = false
	return nil
}

// IsEnabled проверяет, включен ли Kill Switch
func (ks *KillSwitch) IsEnabled() bool {
	return ks.enabled
}