package security

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

// DNSProtection управляет защитой от DNS утечек
type DNSProtection struct {
	originalDNS []string
	enabled     bool
}

// NewDNSProtection создает новый объект защиты DNS
func NewDNSProtection() *DNSProtection {
	return &DNSProtection{
		originalDNS: []string{},
		enabled:     false,
	}
}

// Enable включает защиту от DNS утечек
func (dp *DNSProtection) Enable(vpnDNS []string) error {
	if runtime.GOOS != "windows" {
		return dp.enableUnix(vpnDNS)
	}
	return dp.enableWindows(vpnDNS)
}

// Disable отключает защиту и восстанавливает оригинальные DNS
func (dp *DNSProtection) Disable() error {
	if runtime.GOOS != "windows" {
		return dp.disableUnix()
	}
	return dp.disableWindows()
}

// enableWindows включает защиту DNS на Windows
func (dp *DNSProtection) enableWindows(vpnDNS []string) error {
	// Получаем список активных сетевых интерфейсов
	cmd := exec.Command("netsh", "interface", "show", "interface")
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("ошибка получения списка интерфейсов: %w", err)
	}

	// Парсим вывод для получения имен интерфейсов
	lines := strings.Split(string(output), "\n")
	var activeInterfaces []string
	
	for _, line := range lines {
		if strings.Contains(line, "Connected") || strings.Contains(line, "Подключено") {
			fields := strings.Fields(line)
			if len(fields) >= 4 {
				// Имя интерфейса обычно последнее поле
				interfaceName := strings.Join(fields[3:], " ")
				activeInterfaces = append(activeInterfaces, interfaceName)
			}
		}
	}

	// Сохраняем текущие DNS для каждого интерфейса
	for _, iface := range activeInterfaces {
		cmd = exec.Command("netsh", "interface", "ip", "show", "dns", iface)
		output, err := cmd.Output()
		if err == nil {
			dp.originalDNS = append(dp.originalDNS, string(output))
		}
	}

	// Устанавливаем VPN DNS для всех активных интерфейсов
	for _, iface := range activeInterfaces {
		// Устанавливаем первичный DNS
		if len(vpnDNS) > 0 {
			cmd = exec.Command("netsh", "interface", "ip", "set", "dns",
				iface, "static", vpnDNS[0])
			if err := cmd.Run(); err != nil {
				fmt.Printf("Предупреждение: не удалось установить DNS для %s: %v\n", iface, err)
			}
		}

		// Добавляем вторичные DNS
		for i := 1; i < len(vpnDNS); i++ {
			cmd = exec.Command("netsh", "interface", "ip", "add", "dns",
				iface, vpnDNS[i], fmt.Sprintf("index=%d", i+1))
			cmd.Run() // Игнорируем ошибки для дополнительных DNS
		}
	}

	// Очищаем DNS кэш
	exec.Command("ipconfig", "/flushdns").Run()

	dp.enabled = true
	return nil
}

// disableWindows отключает защиту DNS на Windows
func (dp *DNSProtection) disableWindows() error {
	if !dp.enabled {
		return nil
	}

	// Получаем список активных интерфейсов
	cmd := exec.Command("netsh", "interface", "show", "interface")
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("ошибка получения списка интерфейсов: %w", err)
	}

	lines := strings.Split(string(output), "\n")
	var activeInterfaces []string
	
	for _, line := range lines {
		if strings.Contains(line, "Connected") || strings.Contains(line, "Подключено") {
			fields := strings.Fields(line)
			if len(fields) >= 4 {
				interfaceName := strings.Join(fields[3:], " ")
				activeInterfaces = append(activeInterfaces, interfaceName)
			}
		}
	}

	// Восстанавливаем автоматическое получение DNS
	for _, iface := range activeInterfaces {
		cmd = exec.Command("netsh", "interface", "ip", "set", "dns",
			iface, "dhcp")
		cmd.Run() // Игнорируем ошибки
	}

	// Очищаем DNS кэш
	exec.Command("ipconfig", "/flushdns").Run()

	dp.enabled = false
	dp.originalDNS = []string{}
	return nil
}

// enableUnix включает защиту DNS на Unix системах
func (dp *DNSProtection) enableUnix(vpnDNS []string) error {
	// Сохраняем оригинальный resolv.conf
	cmd := exec.Command("cat", "/etc/resolv.conf")
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("ошибка чтения resolv.conf: %w", err)
	}
	dp.originalDNS = strings.Split(string(output), "\n")

	// Создаем новый resolv.conf с VPN DNS
	var newResolv strings.Builder
	for _, dns := range vpnDNS {
		newResolv.WriteString(fmt.Sprintf("nameserver %s\n", dns))
	}

	// Записываем новый resolv.conf
	cmd = exec.Command("sh", "-c", fmt.Sprintf("echo '%s' > /etc/resolv.conf", newResolv.String()))
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ошибка записи resolv.conf: %w", err)
	}

	dp.enabled = true
	return nil
}

// disableUnix отключает защиту DNS на Unix системах
func (dp *DNSProtection) disableUnix() error {
	if !dp.enabled || len(dp.originalDNS) == 0 {
		return nil
	}

	// Восстанавливаем оригинальный resolv.conf
	content := strings.Join(dp.originalDNS, "\n")
	cmd := exec.Command("sh", "-c", fmt.Sprintf("echo '%s' > /etc/resolv.conf", content))
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ошибка восстановления resolv.conf: %w", err)
	}

	dp.enabled = false
	dp.originalDNS = []string{}
	return nil
}

// IsEnabled проверяет, включена ли защита DNS
func (dp *DNSProtection) IsEnabled() bool {
	return dp.enabled
}

// GetCurrentDNS возвращает текущие DNS серверы
func GetCurrentDNS() ([]string, error) {
	if runtime.GOOS == "windows" {
		return getCurrentDNSWindows()
	}
	return getCurrentDNSUnix()
}

func getCurrentDNSWindows() ([]string, error) {
	cmd := exec.Command("nslookup", "localhost")
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var dnsServers []string
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		if strings.Contains(line, "Address:") {
			parts := strings.Split(line, ":")
			if len(parts) >= 2 {
				dns := strings.TrimSpace(parts[1])
				if dns != "" && dns != "127.0.0.1" {
					dnsServers = append(dnsServers, dns)
				}
			}
		}
	}

	return dnsServers, nil
}

func getCurrentDNSUnix() ([]string, error) {
	cmd := exec.Command("cat", "/etc/resolv.conf")
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var dnsServers []string
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		if strings.HasPrefix(strings.TrimSpace(line), "nameserver") {
			parts := strings.Fields(line)
			if len(parts) >= 2 {
				dnsServers = append(dnsServers, parts[1])
			}
		}
	}

	return dnsServers, nil
}
