package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"
)

func PortIsAvailable(host string, port int) (bool, error) {
	addr := &net.TCPAddr{
		IP:   net.ParseIP(host),
		Port: port,
	}
	listener, err := net.Listen("tcp", addr.String())
	if err != nil {
		if err.Error() == "listen tcp : address already in use" ||
			os.Getenv("GODEBUG") == "netdns=off" &&
				strings.Contains(err.Error(), "address already in use") {
			listener.Close()
			return false, nil
		} else {
			return false, err
		}
	}
	listener.Close()
	return true, nil
}

type proxyResult struct {
	Domain     string
	ReverseUrl *url.URL
}

func main() {
	var e string
	flag.StringVar(&e, "e", "", "domain|ip:port,domain|ip:port,....")
	flag.Parse()

	pattern := `^([\w\.]+)\|(((https|http|ftp|rtsp|mms)?:\/\/)[^\s]+)$`
	re := regexp.MustCompile(pattern)

	proxyMap := make(map[string]proxyResult)

	for _, arg := range strings.Split(e, ",") {
		if re.MatchString(strings.TrimSpace(arg)) {
			domain := re.ReplaceAllString(arg, "$1")
			reverseUrl := re.ReplaceAllString(arg, "$2")
			fmt.Printf("domain proxy => %s -> %s\n", domain, reverseUrl)

			remote, err := url.Parse(reverseUrl)
			if err != nil {
				panic(err)
				return
			}

			proxyMap[domain] = proxyResult{Domain: domain, ReverseUrl: remote}
		} else {
			fmt.Println("The args parameter format is incorrect")
			return
		}
	}

	port := 80
	isAvailable, err := PortIsAvailable("localhost", port)
	if err != nil {
		fmt.Printf("An error occurred while checking the port availability: %v\n", err)
		return
	} else if !isAvailable {
		fmt.Printf("The port %d is already in use.\n", port)
		return
	}

	for _, x := range proxyMap {
		hostEntry := HostEntry{
			IP:   "127.0.0.1",
			Host: x.Domain,
		}
		err2 := addOrUpdateHostEntry(hostEntry)
		if err2 != nil {
			fmt.Printf("Failed to update the hosts file: %v\n", err2)
			return
		} else {
			fmt.Println("Host entry added successfully.")
		}
	}

	var h = &baseHandle{proxyMap}
	http.Handle("/", h)
	server := &http.Server{
		Addr:    ":" + strconv.Itoa(port),
		Handler: h,
	}

	// start http server
	go func() {
		log.Fatal(server.ListenAndServe())
	}()

	// signal notify
	done := make(chan os.Signal, 1)
	signal.Notify(done, syscall.SIGINT, syscall.SIGTERM)

	<-done

	deleteHosts(proxyMap)

	// timeout
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := server.Shutdown(ctx); err != nil {
		log.Fatal("Server Shutdown:", err)
	}

	log.Println("Server exiting")
}

func deleteHosts(proxyMap map[string]proxyResult) {
	for _, x := range proxyMap {
		hostEntry := HostEntry{
			IP:   "127.0.0.1",
			Host: x.Domain,
		}
		err2 := deleteHostsEntry(hostEntry)
		if err2 != nil {
			fmt.Printf("Failed to delete the hosts file: %v\n", err2)
		} else {
			fmt.Println("Host entry delete successfully.")
		}
	}
}