package main

import (
	"context"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Status struct {
	Connected  bool   `json:"connected"`
	Config     string `json:"config"`
	IP         string `json:"ip"`
	Handshake  string `json:"handshake"`
	Ping       string `json:"ping"`
	PacketLoss int    `json:"packetLoss"`
	Error      string `json:"error"`
}

func (s Status) Line() string {
	if !s.Connected {
		if s.Error != "" {
			return "Disconnected | " + s.Error
		}
		return "Disconnected"
	}
	return fmt.Sprintf("Connected to %s | Last handshake %ss | Ping %sms | Packet loss: %d%%", s.IP, s.Handshake, s.Ping, s.PacketLoss)
}

type Engine struct {
	mu          sync.Mutex
	dir         string
	profile     Profile
	conf        string
	cmd         *exec.Cmd
	childDone   chan error
	routes      [][]string
	state       Status
	stopMonitor chan struct{}
	loss        lossWindow
}

func bin(name string) string {
	if runtime.GOOS == "windows" {
		name += ".exe"
	}
	exe, _ := os.Executable()
	for _, p := range []string{filepath.Join(filepath.Dir(exe), "bin", name), filepath.Join(filepath.Dir(exe), name)} {
		if st, e := os.Stat(p); e == nil && !st.IsDir() {
			return p
		}
	}
	if runtime.GOOS == "windows" && (name == "wireguard.exe" || name == "wg.exe") {
		p := filepath.Join(os.Getenv("ProgramFiles"), "WireGuard", name)
		if _, e := os.Stat(p); e == nil {
			return p
		}
	}
	return name
}
func command(name string, args ...string) error {
	c := exec.Command(bin(name), args...)
	c.Env = commandEnv()
	b, e := c.CombinedOutput()
	if e != nil {
		return fmt.Errorf("%s: %w: %s", name, e, strings.TrimSpace(string(b)))
	}
	return nil
}
func commandEnv() []string {
	exe, _ := os.Executable()
	return append(os.Environ(), "NO_COLOR=true", "PATH="+filepath.Join(filepath.Dir(exe), "bin")+string(os.PathListSeparator)+os.Getenv("PATH"))
}
func (e *Engine) activeID() string {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.state.Connected {
		return e.profile.ID
	}
	return ""
}
func (e *Engine) status() Status { e.mu.Lock(); defer e.mu.Unlock(); return e.state }
func (e *Engine) start(p Profile) (err error) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.state.Connected {
		return errors.New("disconnect current profile first")
	}
	if err = validateProfile(p); err != nil {
		return err
	}
	if !admin() {
		return errors.New("WireGuard needs root / Administrator; Termux also needs a WireGuard-enabled Android kernel")
	}
	probe := exec.Command(bin("wg"), "show", "wswg0")
	probe.Env = commandEnv()
	if probe.Run() == nil {
		return errors.New("interface wswg0 already exists; disconnect it before starting")
	}
	e.loss = lossWindow{}
	e.state = Status{Handshake: "—", Ping: "—"}
	e.profile = p
	defer func() {
		if err != nil {
			e.cleanup()
			e.state = Status{Error: err.Error()}
		}
	}()
	w, _ := parseWG(p.WireGuard)
	endpoint := ""
	if p.Tunnel.Enabled {
		host := p.Tunnel.Host
		if host == "" {
			host = w.Host
		}
		ips, x := net.LookupIP(host)
		if x != nil || len(ips) == 0 {
			return fmt.Errorf("resolve wstunnel host %s: %v", host, x)
		}
		var ip net.IP
		for _, v := range ips {
			if v.To4() != nil {
				ip = v
				break
			}
		}
		if ip == nil {
			ip = ips[0]
		}
		if err = e.pinRoute(ip.String()); err != nil {
			return err
		}
		u, x := net.ListenPacket("udp4", "127.0.0.1:0")
		if x != nil {
			return x
		}
		addr := u.LocalAddr().String()
		u.Close()
		_, port, _ := net.SplitHostPort(addr)
		endpoint = addr
		scheme := "ws"
		if p.Tunnel.HTTPS {
			scheme = "wss"
		}
		// The remote UDP destination is resolved by the server. localhost matches a WireGuard server on that host.
		args := []string{"client", "-L", "udp://127.0.0.1:" + port + ":127.0.0.1:" + strconv.Itoa(w.Port) + "?timeout_sec=0"}
		if net.ParseIP(host) == nil {
			if p.Tunnel.HTTPS {
				args = append(args, "--tls-sni-override", host)
			}
			args = append(args, "--http-headers", "Host: "+net.JoinHostPort(host, strconv.Itoa(p.Tunnel.Port)))
		}
		args = append(args, scheme+"://"+net.JoinHostPort(ip.String(), strconv.Itoa(p.Tunnel.Port)))
		if err = e.startTunnel(args); err != nil {
			return err
		}
	}
	e.conf = filepath.Join(e.dir, "wswg0.conf")
	if err = os.WriteFile(e.conf, []byte(renderWG(p, endpoint)), 0600); err != nil {
		return err
	}
	if runtime.GOOS == "windows" {
		err = command("wireguard", "/installtunnelservice", e.conf)
	} else {
		err = command("wg-quick", "up", e.conf)
	}
	if err != nil { // wg-quick rolls back partially configured interfaces itself.
		if runtime.GOOS == "windows" {
			_ = command("wireguard", "/uninstalltunnelservice", "wswg0")
		}
		return err
	}
	e.state = Status{Connected: true, Config: p.Name, IP: w.Host, Handshake: "—", Ping: "—"}
	e.stopMonitor = make(chan struct{})
	go e.monitor(e.stopMonitor)
	return nil
}
func (e *Engine) startTunnel(args []string) error {
	for attempt := 0; attempt < 2; attempt++ {
		log, x := os.OpenFile(filepath.Join(e.dir, "wstunnel.log"), os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if x != nil {
			return x
		}
		c := exec.Command(bin("wstunnel"), args...)
		c.Env = commandEnv()
		c.Stdout = log
		c.Stderr = log
		x = c.Start()
		if x == nil {
			done := make(chan error, 1)
			go func() { done <- c.Wait(); log.Close() }()
			select {
			case x = <-done:
				if x == nil {
					x = errors.New("wstunnel exited")
				}
			case <-time.After(1200 * time.Millisecond):
				e.cmd = c
				e.childDone = done
				return nil
			}
		} else {
			log.Close()
		}
		if runtime.GOOS != "windows" || attempt != 0 || !missingRuntime(x) {
			return fmt.Errorf("wstunnel failed: %w; see %s", x, filepath.Join(e.dir, "wstunnel.log"))
		}
		if x = installRuntime(); x != nil {
			return x
		}
	}
	return errors.New("wstunnel failed after redistributable installation")
}
func missingRuntime(err error) bool {
	s := strings.ToLower(err.Error())
	return strings.Contains(s, "0xc0000135") || strings.Contains(s, "0xc000007b") || strings.Contains(s, "3221225781") || strings.Contains(s, "3221225595")
}
func installRuntime() error {
	arch := runtime.GOARCH
	if arch == "amd64" {
		arch = "x64"
	}
	if arch == "386" {
		arch = "x86"
	}
	exe, _ := os.Executable()
	p := filepath.Join(filepath.Dir(exe), "redist", "vc_redist."+arch+".exe")
	if _, e := os.Stat(p); e != nil {
		return fmt.Errorf("bundled Visual C++ redistributable: %w", e)
	}
	e := exec.Command(p, "/install", "/quiet", "/norestart").Run()
	if x, ok := e.(*exec.ExitError); ok && (x.ExitCode() == 3010 || x.ExitCode() == 1638) {
		return nil
	}
	return e
}
func (e *Engine) cleanup() {
	if e.cmd != nil {
		_ = e.cmd.Process.Kill()
		<-e.childDone
		e.cmd = nil
		e.childDone = nil
	}
	for i := len(e.routes) - 1; i >= 0; i-- {
		r := e.routes[i]
		_ = command(r[0], r[1:]...)
	}
	e.routes = nil
	if e.conf != "" {
		os.Remove(e.conf)
		e.conf = ""
	}
}
func (e *Engine) stop() {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.stopMonitor != nil {
		close(e.stopMonitor)
		e.stopMonitor = nil
	}
	if e.state.Connected {
		var err error
		if runtime.GOOS == "windows" {
			err = command("wireguard", "/uninstalltunnelservice", "wswg0")
		} else {
			err = command("wg-quick", "down", e.conf)
		}
		if err != nil {
			e.state.Error = err.Error()
			return
		}
	}
	e.cleanup()
	e.state = Status{}
}
func (e *Engine) monitor(stop <-chan struct{}) {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-stop:
			return
		case <-ticker.C:
			e.mu.Lock()
			if !e.state.Connected {
				e.mu.Unlock()
				return
			}
			if e.childDone != nil {
				select {
				case err := <-e.childDone:
					e.cmd = nil
					e.childDone = nil
					e.mu.Unlock()
					e.stop()
					e.mu.Lock()
					e.state.Error = fmt.Sprintf("wstunnel exited: %v; see wstunnel.log", err)
					e.mu.Unlock()
					return
				default:
				}
			}
			c := exec.Command(bin("wg"), "show", "wswg0", "latest-handshakes")
			c.Env = commandEnv()
			out, err := c.Output()
			if err == nil {
				latest := int64(0)
				for _, line := range strings.Split(string(out), "\n") {
					f := strings.Fields(line)
					if len(f) == 2 {
						v, _ := strconv.ParseInt(f[1], 10, 64)
						if v > latest {
							latest = v
						}
					}
				}
				if latest > 0 {
					e.state.Handshake = strconv.FormatInt(time.Now().Unix()-latest, 10)
				}
			}
			host := e.state.IP
			e.mu.Unlock()
			ping, success := pingHost(host)
			e.mu.Lock()
			if e.state.Connected && e.stopMonitor == stop {
				e.state.Ping = ping
				e.state.PacketLoss = e.loss.add(success)
			}
			e.mu.Unlock()
		}
	}
}
func pingHost(host string) (string, bool) {
	ctx, cancel := context.WithTimeout(context.Background(), 950*time.Millisecond)
	defer cancel()
	args := []string{"-n", "-c", "1", "-W", "1", host}
	if runtime.GOOS == "windows" {
		args = []string{"-n", "1", "-w", "1000", host}
	}
	out, err := exec.CommandContext(ctx, "ping", args...).Output()
	if err != nil {
		return "—", false
	}
	s := strings.ToLower(string(out))
	// Windows ping uses the system OEM encoding. TTL remains ASCII even when
	// the time label and the unit are localized.
	if runtime.GOOS == "windows" {
		for _, line := range strings.Split(s, "\n") {
			f := strings.Fields(line)
			for i, part := range f {
				if i > 0 && strings.HasPrefix(part, "ttl=") {
					value := f[i-1]
					if j := strings.LastIndexAny(value, "=<"); j >= 0 {
						value = value[j+1:]
						k := 0
						for k < len(value) && ((value[k] >= '0' && value[k] <= '9') || value[k] == '.' || value[k] == ',') {
							k++
						}
						if k > 0 {
							return value[:k], true
						}
					}
				}
			}
		}
	}
	for _, key := range []string{"time=", "time<", "время="} {
		i := strings.Index(s, key)
		if i >= 0 {
			tail := s[i+len(key):]
			end := 0
			for end < len(tail) && ((tail[end] >= '0' && tail[end] <= '9') || tail[end] == '.' || tail[end] == ',') {
				end++
			}
			if end > 0 {
				return tail[:end], true
			}
		}
	}
	return "—", true
}

// lossWindow stores the last 100 ICMP attempts; an unanswered probe is a loss.
type lossWindow struct {
	results           [100]bool
	next, count, lost int
}

func (w *lossWindow) add(success bool) int {
	if w.count == len(w.results) {
		if !w.results[w.next] {
			w.lost--
		}
	} else {
		w.count++
	}
	w.results[w.next] = success
	if !success {
		w.lost++
	}
	w.next = (w.next + 1) % len(w.results)
	return (w.lost*100 + w.count/2) / w.count
}