// Copyright 2017-2025 WIT.COM Inc. All rights reserved.
// Use of this source code is governed by the GPL 3.0

package main

import (
	"bufio"
	"embed"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	sync "sync"
	"time"

	"github.com/svent/go-nbreader"

	"github.com/google/uuid"
	"go.wit.com/dev/alexflint/arg"
	"go.wit.com/gui"
	"go.wit.com/log"
	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
)

var VERSION string
var BUILDTIME string

//go:embed resources/*
var resources embed.FS

func main() {
	var pp *arg.Parser
	gui.InitArg()
	pp = arg.MustParse(&argv)

	if pp == nil {
		pp.WriteHelp(os.Stdout)
		os.Exit(0)
	}
	log.Info("tmp hack", uuid.New().String())

	me = new(gusconf)
	me.pollDelay = 10 * time.Second

	if argv.UseME {
		connectME()
		okExit("")
	}

	me.portmaps = ConfigLoad()
	me.events = EventLoad()

	if me.portmaps == nil {
		me.portmaps = NewPortmaps()
		p := new(Portmap)
		p.Dest = "testing:323"
		me.portmaps.Append(p)
	}

	if argv.Daemon {
		// turn off timestamps for STDOUT (systemd adds them)
		log.DaemonMode(true)
		startGus()
		startHTTP()
		os.Exit(0)
	}

	if gui.NoGui() {
		startGus()
		startHTTP()
		os.Exit(0)
	}

	startGus()
	// go NewWatchdog()
	go startHTTP()
	doGui()
}

func startGus() {
	all := me.portmaps.All()
	for all.Scan() {
		pm := all.Next()
		if !pm.Enabled {
			continue
		}
		log.Info("portmap enabled for port", pm.Localport, "to", pm.Dest)
		go gus3000(pm)
	}
}

// func doME(pm *Portmap, gus listener.Accept) {
func doME(pm *Portmap, gus net.Listener) {
	localport := int(pm.Localport)
	where := pm.Dest

	/*
		// Listen on local port 3000
		s := fmt.Sprintf("0.0.0.0:%d", port)
		listener, err := net.Listen("tcp", s)
		if err != nil {
			log.Fatalf("Failed to listen on %s: %v", s, err)
		}
		defer listener.Close()
		log.Info("Listening on ", s)
	*/

	for {
		// Accept incoming connection
		src, err := gus.Accept()
		if err != nil {
			log.Printf("Failed to accept client connection: %v\n", err)
			pm.Enabled = false
			return
		}

		// make a new event from this new connection
		log.Printf("/me Connected on port %d from client: %s to where = %s\n", localport, src.RemoteAddr(), where)
		fmt.Fprintln(src, "/me hello")

		reader := bufio.NewReader(src)

		// Read one line
		line, err := reader.ReadString('\n')
		if err != nil {
			log.Info("gus src", src.RemoteAddr(), "read error:", err)
			continue
		}
		if !strings.HasPrefix(line, "/me") {
			log.Printf("gus Received %d invalid bytes\n", len(line))
			continue
		}
		parts := strings.Fields(line)
		if len(parts) != 3 {
			continue
		}
		if parts[0] != "/me" {
			continue
		}
		if parts[1] != "hostname" {
			continue
		}
		hostname := parts[2]
		msg := fmt.Sprintf("got hostname %s for %s", hostname, src.RemoteAddr())
		log.Info("gus:", msg)
		fmt.Fprintln(src, msg)

		if hostname == "framebook.wit.com" {
			// Handle the connection in a separate goroutine
			log.Info("RUNNING DIAL")
			dest, err := net.Dial("tcp", where)
			if err != nil {
				log.Printf("Failed to connect to %s %v", where, err)
				continue
			}
			defer dest.Close()

			log.Info("IOCOPY START")
			ioCopy(src, dest)
			log.Info("IOCOPY END")
		}
	}
}

func gus3000(pm *Portmap) {
	port := int(pm.Localport)
	connect := pm.Dest

	// Listen on local port 3000
	s := fmt.Sprintf("0.0.0.0:%d", port)
	listener, err := net.Listen("tcp", s)
	if err != nil {
		log.Fatalf("Failed to listen on %s: %v", s, err)
	}
	defer listener.Close()
	log.Info("Listening on ", s)

	if pm.UseME {
		doME(pm, listener)
		return
	}

	for {
		// Accept incoming connection
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept client connection: %v", err)
			continue
		}
		// log.Printf("Client connected: %s", clientConn.RemoteAddr())

		// Handle the connection in a separate goroutine
		go handleConnection(clientConn, connect, port)
	}
}

func handleConnection(clientConn net.Conn, where string, localport int) {
	defer clientConn.Close()

	// Connect to the target server
	// targetConn, err := net.Dial("tcp", "go.wit.com:80")
	targetConn, err := net.Dial("tcp", where)
	if err != nil {
		log.Printf("Failed to connect to %s %v", where, err)
		return
	}
	defer targetConn.Close()

	// make a new event from this new connection
	log.Printf("Connected on port %d from client: %s to where = %s\n", localport, clientConn.RemoteAddr(), where)
	e := new(Event)
	e.Etype = GusEventType_Connect
	e.LocalPort = int64(localport)
	e.Sock = new(GusSocket)
	e.Sock.SrcIp = fmt.Sprintf("%s", clientConn.RemoteAddr())
	e.Sock.DestIp = where
	e.Ctime = timestamppb.New(time.Now())
	me.events.Append(e)
	me.eventsChanged = true

	// Bidirectional copy of data
	// go io.Copy(targetConn, clientConn) // Client -> Target
	// io.Copy(clientConn, targetConn)    // Target -> Client
	ioCopy(clientConn, targetConn)

	// if the socket closes, record the close time
	e.Etime = timestamppb.New(time.Now())
	me.eventsChanged = true
	log.Printf("Connection closed on port %d from client: %s to where = %s\n", localport, clientConn.RemoteAddr(), where)
}

func enablePort(port int, dest string) {
	var found bool
	all := me.portmaps.All()
	for all.Scan() {
		pm := all.Next()
		if int(pm.Localport) == port {
			found = true
			log.Info("Found port!", port)
			if pm.Enabled {
				log.Info("portmap already enabled for", pm.Localport, "to", pm.Dest)
			} else {
				log.Info("portmap not enabled for", pm.Localport, "to", pm.Dest)
			}
		}
		// go gus3000(pm)
	}
	if !found {
		log.Info("Did not find port =", port)
	}
}

// should cleanly close both
func ioCopy(clientConn, targetConn net.Conn) {
	defer clientConn.Close()
	defer targetConn.Close()

	/*
		if !argv.UseME {
			// Drain existing data
			drainConnection(clientConn)
			drainConnection(targetConn)
		}
	*/

	var wg sync.WaitGroup
	wg.Add(2)

	// Copy data in both directions
	go func() {
		defer wg.Done()
		io.Copy(targetConn, clientConn) // Client -> Target
		targetConn.Close()              // Ensure closure on EOF/error
	}()

	go func() {
		defer wg.Done()
		io.Copy(clientConn, targetConn) // Target -> Client
		clientConn.Close()              // Ensure closure on EOF/error
	}()

	wg.Wait() // Wait for both copies to complete
}

func connectME() {
	hostname, err := os.Hostname()
	if err != nil {
		fmt.Println("Error:", err)
		return
	}
	log.Println("Hostname:", hostname)
	if fqdn, err := getFQDN(hostname); err == nil {
		hostname = fqdn
	} else {
		log.Printf("Your hostname (%s) is not in DNS correctly: %v\n", hostname, err)
	}

	localport := "25910"
	where := argv.URL
	dest, err := net.Dial("tcp", where)
	if err != nil {
		log.Printf("Failed to connect to %s %v\n", where, err)
		return
	}
	defer dest.Close()

	// make a new event from this new connection
	log.Printf("Connected on port %s from client: %s to where = %s\n", localport, "tbd", where)

	reader := bufio.NewReader(dest)

	// Read one line
	line, err := reader.ReadString('\n')
	if err != nil {
		log.Info("Error reading line:", err)
		return
	}
	log.Info("gus Received:", line)

	fmt.Fprintf(dest, "/me hostname %s\n", hostname)

	// Read one line
	line, err = reader.ReadString('\n')
	if err != nil {
		log.Info("Error reading line:", err)
		return
	}
	log.Info("gus Received:", line)

	// Listen on local port
	s := fmt.Sprintf("0.0.0.0:%s", localport)
	listener, err := net.Listen("tcp", s)
	if err != nil {
		log.Fatalf("Failed to listen on %s: %v", s, err)
	}
	defer listener.Close()
	log.Info("Listening on ", s)
	log.Printf("Try: remmina -c spice://localhost:%s\n", localport)

	// Accept incoming connection
	src, err := listener.Accept()
	if err != nil {
		log.Printf("Failed to accept client connection: %v", err)
		return
	}
	log.Printf("Client connected: %s\n", src.RemoteAddr())
	ioCopy(src, dest)
}

func getFQDN(hostname string) (string, error) {
	// Perform a reverse lookup to get the FQDN
	addrs, err := net.LookupAddr(hostname)
	if err != nil {
		return "", fmt.Errorf("failed to resolve FQDN: %w", err)
	}

	// Choose the first valid result
	for _, addr := range addrs {
		if strings.HasSuffix(addr, ".") { // FQDNs often end with a dot
			return strings.TrimSuffix(addr, "."), nil
		}
		return addr, nil
	}

	return "", fmt.Errorf("no FQDN found for hostname: %s", hostname)
}

// drainConnection reads and discards any pending data before copying
func drainConnection(conn net.Conn) {
	// Create a non-blocking reader with a timeout of 1 second
	reader := nbreader.NewNBReader(conn, 1024, nbreader.Timeout(time.Second))

	buffer := make([]byte, 1024)

	var count int

	for {
		// Attempt to read data from the reader
		n, err := reader.Read(buffer)
		if err != nil {
			log.Info("error was:", err)
			// Handle error, such as timeout or end of stream
			break
		}

		if n > 0 {
			log.Info("chunk size was", n, "buffer =", string(buffer))
			count = 0
			continue
		}

		count += 1
		if count > 3 {
			log.Info("buffer is drained")
			return
		}

		// Process the read data (buffer[:n])
	}
}