package main

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/netip"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/biter777/countries"
	"github.com/google/uuid"
	"github.com/urfave/cli/v2"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	hvpnnode3 "gitea.hbanafa.com/HeshamTB/hvpn-node3"
)

/*
   Can change all therse vars to local func vars
   IPPool can be internal to wglink or other abstraction. As well as
   Port, ifname and CIDR and so on. Use Clie.ctx to get the values
*/

var IPPool hvpnnode3.IPPool
var VPNIPCIDR string
var PrivateKeyPath cli.Path
var InterfaceName string
var WgPort int
var wgLink *hvpnnode3.WGLink

var Country countries.CountryCode
var UUID uuid.UUID

var httpPort int
var TLS_ENABLED bool
var tlsConfig *tls.Config

func main() {


    // TODO: Define error exit codes
    slog.SetLogLoggerLevel(slog.LevelInfo)


    app := createCliApp()
    err := app.Run(os.Args)
    if err != nil {
        slog.Error(err.Error())
        os.Exit(1)
    }
}

func run(ctx *cli.Context) {
    slog.Debug("Starting run()")

    wgLink.StartedAT = time.Now().UTC()
    slog.Info(fmt.Sprintf("Started at %s", wgLink.StartedAT))
    slog.Info(fmt.Sprintf("Country set to %s (%s)", Country.Alpha2(), Country.String()))
    hvpnnode3.StartMonitor(wgLink, *slog.Default())

    apiMux := http.NewServeMux()
    apiMux.HandleFunc("GET /node", hvpnnode3.HandleGetNodeInfo(wgLink))
    apiMux.HandleFunc("GET /peer/{pubkey}", hvpnnode3.HandleGetPeer(wgLink))
    apiMux.HandleFunc("POST /peer", hvpnnode3.HandlePostPeer(wgLink))
    apiMux.HandleFunc("DELETE /peer/{pubkey}", hvpnnode3.HandleDeletePeer(wgLink))
    apiMux.HandleFunc("GET /peers", hvpnnode3.HandleGetPeers(wgLink))
    apiMux.HandleFunc("GET /peer", hvpnnode3.HandleGetCreatePeer(wgLink))

    var handler http.Handler = apiMux
    handler = hvpnnode3.WithCountryCtx(handler, Country)
    handler = hvpnnode3.HttpAuthToken(handler, ctx.String("http-api-key"))
    handler = hvpnnode3.HttpLogHandler2(handler)

    port := fmt.Sprintf("%d", httpPort)
    host := ctx.String("host")
    if host == "" {
        slog.Info("Host is not set. Using 0.0.0.0")
        host = "0.0.0.0"
    }
    hostPort := fmt.Sprintf("%s:%s", host, port)

    slog.Info(fmt.Sprintf("Starting HVPN node on %s", hostPort))
    srv := &http.Server{
        ReadTimeout: 120 * time.Second,
        WriteTimeout: 120 * time.Second,
        IdleTimeout: 120 * time.Second,
        Handler: handler,
        Addr: hostPort,
        TLSConfig: tlsConfig,
    }

    defer wgLink.Close()
    if TLS_ENABLED {
        slog.Debug("Running with TLS")
        slog.Info("Running TLS or mTLS with a reverse proxy is recommended")
        slog.Warn(srv.ListenAndServeTLS(ctx.Path("cert"), ctx.Path("cert-private-key")).Error())
    } else {
        slog.Warn(srv.ListenAndServe().Error())
    }
}

func createCliApp() *cli.App {
    app := cli.NewApp()
    app.Name = os.Args[0]
    app.HideHelpCommand = true
    app.Usage = "HVPN node API server"
    app.Args = true

    author1 := cli.Author{
        Name: "Hesham T. Banafa",
        Email: "hishaminv@gmail.com"}
    app.Authors = append(app.Authors, &author1)
    
    logLevel := cli.StringFlag{
        Name: "log-level",
        Value: "INFO",
        EnvVars: []string{"LOG_LEVEL"},
        Action: func(ctx *cli.Context, s string) error {
            lvl := new(slog.LevelVar)
            err := lvl.UnmarshalText([]byte(s))
            if err != nil {
                slog.Debug("Error on unmarshal log level flag")
                return err
            }
            slog.SetLogLoggerLevel(lvl.Level())
            return nil 
        },
    }
    app.Flags = append(app.Flags, &logLevel)

    privateKeyFileFlag := cli.PathFlag{
        Name: "private-key",
        Usage: "Path to file with private key",
        Destination: &PrivateKeyPath,
    }
    app.Flags = append(app.Flags, &privateKeyFileFlag)

    vpnIpCIDR := cli.StringFlag{
        Name: "cidr",
        Usage: "The network subnet used for the internal IP Pool",
        Value: "10.42.0.0/16",
        Aliases: []string{"n"},
        Destination: &VPNIPCIDR,
    }
    app.Flags = append(app.Flags, &vpnIpCIDR)

    wgInterfaceName := cli.StringFlag{
        Name: "interface",
        Usage: "Name of the Wireguard interface to be created and managed",
        Value: "hvpn0",
        Aliases: []string{"i"},
        Destination: &InterfaceName,
    }
    app.Flags = append(app.Flags, &wgInterfaceName)

    wgEndpoint := cli.StringFlag{
        Name: "endpoint",
        Usage: "Wireguard endpoint domain or address without the port",
        Value: "domain.name.notset",
    }
    app.Flags = append(app.Flags, &wgEndpoint)

    wgPort := cli.IntFlag{
        Name: "port",
        Usage: "UDP Port for wireguard device",
        Value: 6416,
        Aliases: []string{"p"},
        Destination: &WgPort,
    }
    app.Flags = append(app.Flags, &wgPort)

    httpListenAddr := cli.StringFlag{
        Name: "http-host",
        Usage: "IP address to listen on for HTTP API requests",
        Value: "0.0.0.0",
        Action: func(ctx *cli.Context, s string) error {
            _, err := netip.ParseAddr(s)
            if err != nil {
                return errors.New(fmt.Sprintf("Can not parse %s as a network IP", s))
            }
            return nil
        },
    }
    app.Flags = append(app.Flags, &httpListenAddr)

    httpPort := cli.IntFlag{
        Name: "http-port",
        Usage: "TCP Port for HTTP API",
        Value: 8080,
        Destination: &httpPort,
    }
    app.Flags = append(app.Flags, &httpPort)

    apiSecret := cli.StringFlag{
        Name: "http-api-key",
        Usage: "Secure endpoints with this key; 'Authorization: Bearer <key>' HTTP Header",
    }
    app.Flags = append(app.Flags, &apiSecret)

    countryFlag := cli.StringFlag{
        Name: "country",
        Usage: "The contry code for this VPN instance. Accepts ISO 3166 codes",
        Value: "SA",
        Action: func(ctx *cli.Context, s string) error {
            c := countries.ByName(s) 
            if c == countries.Unknown {
                return errors.New(fmt.Sprintf("Country code %s is unknown", s))
            }
            return nil
        },
        EnvVars: []string{ "HVPN_COUNTRY" },
    }
    app.Flags = append(app.Flags, &countryFlag)


    /* TLS Flags */

    TLS := cli.BoolFlag{
        Name: "enable-tls",
        Aliases: []string{ "tls" },
        Value: false,
        Category: "\rTLS:",
        Destination: &TLS_ENABLED,
        Action: func(ctx *cli.Context, b bool) error {
            tlsConf, err := setupTLS(ctx)
            if err != nil {
                return err
            }
            tlsConfig = tlsConf
            return nil
        },
    }
    app.Flags = append(app.Flags, &TLS)

    mTLSClientCerts := cli.PathFlag{
        Name: "client-certs",
        Aliases: []string{"ca"},
        Usage: "Clients x509 file with single or many certificates",
        Category: "\rTLS:",
    }
    app.Flags = append(app.Flags, &mTLSClientCerts)

    TLSCert := cli.PathFlag{
        Name: "cert",
        Usage: "Server x509 certificate file",
        Category: "\rTLS:",
    }
    app.Flags = append(app.Flags, &TLSCert)

    TLSCertKey := cli.PathFlag{
        Name: "cert-private-key",
        Usage: "Server x509 certificate private key file",
        Category: "\rTLS:",
    }
    app.Flags = append(app.Flags, &TLSCertKey)


    app.Action = func(ctx *cli.Context) error {
        err := setup(ctx)
        if err != nil {
            return cli.Exit(err, 1)
        }
        run(ctx)
        return nil
    }

    return app
}

func setup(ctx *cli.Context) error {
    slog.Debug("Starting setup()")
        uid := os.Getuid()
    if uid == -1 {
        slog.Warn("Running on windows! whatrr u doing?")
    } else if uid == 0 {
        slog.Warn("Running as root! avoid running as root by setting CAP_NET_ADMIN")
    }
    uuid, err := hvpnnode3.InitNodeUUID()
    if err != nil {
        slog.Error(err.Error())
        os.Exit(-1)
    }
    slog.Info("Node UUID: " + uuid.String())
    UUID = *uuid
    Country = countries.ByName(string(ctx.String("country")))

    var privateKey wgtypes.Key
    createPrivKey := func() error {
        slog.Info("Creating a private key")
        privateKey, err = wgtypes.GeneratePrivateKey()
        if err != nil {
            return err
        }
        slog.Debug(fmt.Sprintf("new public key: %s", privateKey.PublicKey().String()))
        return nil
    }

    if PrivateKeyPath == "" {
        err := createPrivKey()
        if err != nil {
            return err
        }
    } else {
        privKeyFile, err := os.Open(PrivateKeyPath)
        defer privKeyFile.Close()
        if err != nil {
            slog.Error(err.Error())
            slog.Info("Could not open private key file")
            err := createPrivKey()
            if err != nil {
                return err
            }
        } else {
            privateKeyStr := make([]byte, 45)
            n, err := privKeyFile.Read(privateKeyStr)
            if err != nil {
                return err
            }
            if n != 45 {
                slog.Warn("Private key length did not math the expected 45!")
            }
            slog.Debug(fmt.Sprintf("Read %d bytes from keyfile", n))

            privateKey, err = wgtypes.ParseKey(string(privateKeyStr))
            slog.Debug("Keyfile opened for reading")
            if err != nil {
                return err
            }
            slog.Debug("Private key parsed and is correct")
        }
    }

    wg, err := hvpnnode3.InitWGLink(
        InterfaceName, 
        &privateKey,
        WgPort,
        ctx.String("endpoint"),
    )
    if err != nil {
        slog.Warn("Error while initlizing Wireguard netlink and device!")
        slog.Warn("Ensure to run as root or with CAP_NET_ADMIN")
        return err
    }

    wgLink = wg

    slog.Info("Wireguard netlink is setup and running")


    dev, err := wgLink.Device(InterfaceName)
    if err != nil {
        return err
    }
    slog.Info("Initialized Wireguard device using wgctrl")

    slog.Info(
        fmt.Sprintf(
            "Device name: %s, UDP port: %d, Type: %s",
            dev.Name, dev.ListenPort, dev.Type.String(),
        ),
    )

    ipPool, err := hvpnnode3.NewPool(VPNIPCIDR)
    if err != nil {
        slog.Error(fmt.Sprintf("IPPool: %s", err))
        os.Exit(1)
    }
    slog.Debug(fmt.Sprintf("Init ip pool %s", VPNIPCIDR))

    testVip, err := ipPool.Allocate()
    if err != nil {
        slog.Error("main.testVip: ", err)
        os.Exit(1)
    }

    slog.Debug(fmt.Sprintf("IP Pool Test IP: %s", testVip.String()))
    err = ipPool.Free(testVip)
    if err != nil {
        slog.Error("Could not free test Vip from IPPool!", err)
        return err
    }
    slog.Debug("Test IP Freed")

    IPPool = ipPool
    wgLink.IPPool = ipPool

    err = wgLink.SetIP()
    if err != nil {
        return err
    }
    slog.Debug("Assigned IP to Wiregaurd interface")

    //defer wgLink.Close()
    cInput := make(chan struct{})
    go handleStdin(cInput)
    c := make(chan os.Signal, 1)
    signal.Notify(c, os.Interrupt, os.Kill)
    go func() {
        slog.Info("Listening for SIGINT, SIGKILL and user input")
        select {
        case <- c:
            slog.Warn("Recived SIGINT! Closing Wireguard Device")
        case <- cInput:
            slog.Warn("Recieved quit! Closing Wireguard Device")
        }
        wgLink.Close()
        os.Exit(0)
    }()

    err = testWgPeerAdd(wgLink)
    if err != nil {
        return err
    }

    return nil
}

func testWgPeerAdd(wgLink *hvpnnode3.WGLink) error {
    privateKey, err := wgtypes.GeneratePrivateKey()
    if err != nil {
        return err
    }
    publicKey := privateKey.PublicKey()
    _, err = wgLink.AddPeer(publicKey.String())
    if err != nil {
        slog.Error(err.Error())
        return err
    }
    slog.Debug(fmt.Sprintf("Added test peer %v", publicKey.String()))
    
    err = wgLink.DeletePeer(publicKey.String())
    if err != nil {
        return err
    }

    peers, err := wgLink.GetAllPeers()
    if len(peers) != 0 {
        slog.Warn(fmt.Sprintf("Expected 0 peers, got %d", len(peers)))
    }

    slog.Debug("Removed test peer")

    return nil
}

func handleStdin(c chan struct{}) {
    for {
        reader := bufio.NewReader(os.Stdin)
        slog.Info("Enter 'q' or 'exit' to quit")
        in, _ := reader.ReadString('\n')
        in = strings.ReplaceAll(in, "\n", "")
        if in == "q" || in == "exit" {
            c <- struct{}{}
            break
        }
    }
}

func createMTLS(clinetCert, serverCert, serverKey string) (*x509.CertPool, error) {
    clientCert, err := os.ReadFile(clinetCert)
    if err != nil {
        return nil, err
    }
    certPool := x509.NewCertPool()
    ok := certPool.AppendCertsFromPEM(clientCert)
    if !ok {
        return nil, errors.New("mTLS: Could read and parase a single cert from client cert")
    }
    return certPool, nil
}

func setupTLS(ctx *cli.Context) (*tls.Config, error) {
    ca := ctx.Path("client-certs")
    if ca == "" {
        return nil, errors.New("client-certs flag is not set to enable TLS")
    }

    serverCert := ctx.Path("ca")
    if serverCert == "" {
        return nil, errors.New("cert flag is not set to enable TLS")
    }
    
    serverKey := ctx.Path("cert-private-key")
    if serverKey == "" {
        return nil, errors.New("cert-private-key is not set to enable TLS")
    }

    certPool, err := createMTLS(ca, serverCert, serverKey)
    if err != nil {
        return nil, err
    }

    conf := tls.Config{
        ClientAuth: tls.RequireAndVerifyClientCert,
        ClientCAs: certPool,
    }

    return &conf, nil
}