package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/spf13/cobra"
)

var rootCmd = &cobra.Command{
	Use:     "websocket-proxy",
	Version: "0.1.0",
	Short:   "Proxy TCP connections over a websocket",
}

var clientCmd = &cobra.Command{
	Use:   "client [server host]",
	Short: "Act as a client for a websocket-proxy server",
	Args: func(cmd *cobra.Command, args []string) error {
		if len(args) != 1 || args[0] == "" {
			return errors.New("Server host is a required argument")
		}
		if !strings.HasPrefix(args[0], "ws://") && !strings.HasPrefix(args[0], "wss://") {
			return errors.New("Server host format is ws[s]://host[:port]/[path]")
		}
		return nil
	},
	Run: func(cmd *cobra.Command, args []string) {
		listenOn := cmd.Flag("listen").Value.String()

		h := &ClientHandler{
			SocketListenOn:  listenOn,
			WebsocketServer: args[0],
		}

		log.Printf("Serving on %s", listenOn)
		h.Run()
	},
}

var serverCmd = &cobra.Command{
	Use:   "server [next-hop host]",
	Short: "Serve websocket proxy client",
	Args: func(cmd *cobra.Command, args []string) error {
		if len(args) != 1 || args[0] == "" {
			return errors.New("Next-hop host is a required argument")
		}
		return nil
	},
	Run: func(cmd *cobra.Command, args []string) {
		listenOn := cmd.Flag("listen").Value.String()
		log.Printf("Serving on %s", listenOn)

		http.Handle("/", NewServerHandler(args[0]))
		log.Fatal(http.ListenAndServe(listenOn, nil))
	},
}

func main() {
	rootCmd.AddCommand(clientCmd)
	rootCmd.AddCommand(serverCmd)

	clientCmd.Flags().StringP("listen", "l", ":9013", "[address]:port to bind for serving clients")
	serverCmd.Flags().StringP("listen", "l", ":9012", "[address]:port to bind for serving clients")

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}