package portforward

import (
	"context"
	"net"
	"sync"
	"time"

	"github.com/huin/goupnp/dcps/internetgateway1"
	"github.com/veggiedefender/torrent-client/internal/logger"
)

type PortManager struct {
	mu      sync.Mutex
	clients []*internetgateway1.WANIPConnection1
	port    uint16
	opened  bool
}

func NewPortManager() *PortManager {
	return &PortManager{}
}

// OpenPort attempts to open the specified port on all discovered UPnP routers.
func (pm *PortManager) OpenPort(ctx context.Context, port uint16, desc string) error {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if pm.opened && pm.port == port {
		return nil
	}

	logger.Info("SYS", "Discovering UPnP routers for port %d...", port)

	discoverCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()

	discoverDone := make(chan []*internetgateway1.WANIPConnection1, 1)
	go func() {
		clients, _, _ := internetgateway1.NewWANIPConnection1Clients()
		discoverDone <- clients
	}()

	var clients []*internetgateway1.WANIPConnection1
	select {
	case <-discoverCtx.Done():
		logger.Warn("SYS", "UPnP discovery timed out")
		return nil
	case clients = <-discoverDone:
	}

	if len(clients) == 0 {
		logger.Warn("SYS", "No UPnP routers discovered")
		return nil
	}

	localIP, err := getLocalIP()
	if err != nil {
		logger.Warn("SYS", "Could not get local IP for UPnP")
		return nil
	}

	success := false
	for _, client := range clients {
		// TCP Mapping
		errTCP := client.AddPortMapping("", port, "TCP", port, localIP, true, desc, 0)
		// UDP Mapping (for DHT)
		errUDP := client.AddPortMapping("", port, "UDP", port, localIP, true, desc, 0)
		
		if errTCP == nil || errUDP == nil {
			success = true
		}
	}

	if success {
		pm.clients = clients
		pm.port = port
		pm.opened = true
		logger.Info("SYS", "UPnP successfully forwarded port %d", port)
	} else {
		logger.Warn("SYS", "Failed to add UPnP port mapping")
	}

	return nil
}

// ClosePort closes the currently opened port.
func (pm *PortManager) ClosePort() error {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if !pm.opened || len(pm.clients) == 0 {
		return nil
	}

	for _, client := range pm.clients {
		_ = client.DeletePortMapping("", pm.port, "TCP")
		_ = client.DeletePortMapping("", pm.port, "UDP")
	}

	logger.Info("SYS", "UPnP successfully cleared port %d", pm.port)
	pm.opened = false
	return nil
}

// getLocalIP returns the preferred outbound IP of this machine
func getLocalIP() (string, error) {
	conn, err := net.Dial("udp", "8.8.8.8:80")
	if err != nil {
		return "", err
	}
	defer conn.Close()
	localAddr := conn.LocalAddr().(*net.UDPAddr)
	return localAddr.IP.String(), nil
}