package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"flag"

	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
)

const charset = "abcdefghijklmnopqrstuvwxyz0123456789"

func generateSlug(length int) string {
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return string(b)
}

// TunnelManager manages the active tunnels
type TunnelManager struct {
	tunnels map[string]*Tunnel
	mu      sync.RWMutex
}

type Tunnel struct {
	ID        string
	LocalPort uint32
	Conn      *gossh.ServerConn
}

var manager = &TunnelManager{
	tunnels: make(map[string]*Tunnel),
}

// requestedSlugs stores slugs requested by clients before tcpip-forward
var (
	requestedSlugs   = make(map[*gossh.ServerConn]string)
	requestedSlugsMu sync.Mutex
	slugPattern      = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{1,30}[a-z0-9]$`)
)

func (tm *TunnelManager) Register(id string, t *Tunnel) {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	tm.tunnels[id] = t
}

func (tm *TunnelManager) Unregister(id string) {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	delete(tm.tunnels, id)
}

func (tm *TunnelManager) Get(id string) (*Tunnel, bool) {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	t, ok := tm.tunnels[id]
	return t, ok
}

func (tm *TunnelManager) FindSlugByConn(conn *gossh.ServerConn) (string, bool) {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	for slug, t := range tm.tunnels {
		if t.Conn == conn {
			return slug, true
		}
	}
	return "", false
}

// HTTP Proxy to forward requests to the SSH tunnel
func startHttpProxy(port string) {
	log.Printf("Starting HTTP Proxy on %s", port)
	// Serve static files for 404 page (namespaced to avoid collision with user apps)
	http.Handle("/_gw/", http.StripPrefix("/_gw/", http.FileServer(http.Dir("./cmd/server/static/assets"))))

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Extract subdomain to identify the tunnel
		// Format expects: slug.grokway.com or slug.localhost
		host := r.Host
		parts := strings.Split(host, ".")
		// If accessing IP directly or localhost without subdomain
		if len(parts) < 2 || host == "localhost:8080" {
			http.Error(w, "Grokway Server Running", http.StatusOK)
			return
		}
		slug := parts[0]

		log.Printf("Request Host: %s, Extracted Slug: %s", host, slug)

		tunnel, ok := manager.Get(slug)
		if !ok {
			log.Printf("Tunnel not found for slug: %s", slug)
			// Serve 404 page
			w.WriteHeader(http.StatusNotFound)
			http.ServeFile(w, r, "./cmd/server/static/404.html")
			return
		}

		// In a real implementation, we would dial the tunnel
		// using ssh.Conn.OpenChannel("direct-tcpip", ...)
		// But gliderlabs/ssh handles Reverse Port Forwarding differently.
		// It usually binds a listener on the server side.

		// Wait! The standard "Remote Port Forwarding" (-R) means the SSH Server
		// listens on a port ON THE SERVER, and forwards traffic to the client.

		// If GrokwayV1 used -R 80:localhost:8080, then the server was listening on 80.
		// Here we want dynamic subdomains.
		// So we likely want to use a Custom Request or just standard forwarding
		// but mapped to our HTTP proxy.

		// SIMPLE APPROACH (Like Ngrok):
		// 1. Client connects via SSH.
		// 2. Client requests "tcpip-forward" (remote forwarding).
		// 3. Server acknowledges.
		// 4. When HTTP request comes to `slug.grokway.com`:
		//    a. Server accepts request.
		//    b. Server opens a "forwarded-tcpip" channel TO the client.
		//    c. Server proxies the HTTP request body over this channel.

		// For this MVP, let's assume we implement the "forwarded-tcpip" logic.
		// However, gliderlabs/ssh simplifies this if we use its Forwarding support.
		// But typically `gliderlabs/ssh` is for allowing the server to be a jump host.
		// We want to be an HTTP Gateway.

		// Panic recovery
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered from panic in handler: %v", r)
				http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			}
		}()

		// 2. Open channel to client
		destHost := "0.0.0.0"
		destPort := tunnel.LocalPort
		srcHost := "127.0.0.1"
		srcPort := uint32(12345)

		payload := make([]byte, 0)
		payload = append(payload, []byte{0, 0, 0, uint8(len(destHost))}...)
		payload = append(payload, []byte(destHost)...)
		payload = append(payload, byte(destPort>>24), byte(destPort>>16), byte(destPort>>8), byte(destPort))
		payload = append(payload, []byte{0, 0, 0, uint8(len(srcHost))}...)
		payload = append(payload, []byte(srcHost)...)
		payload = append(payload, byte(srcPort>>24), byte(srcPort>>16), byte(srcPort>>8), byte(srcPort))

		ch, reqs, err := tunnel.Conn.OpenChannel("forwarded-tcpip", payload)
		if err != nil {
			log.Printf("Failed to open forwarded-tcpip channel: %s", err)
			return
		}
		defer ch.Close()
		go gossh.DiscardRequests(reqs)

		// Check if it is a WebSocket Upgrade
		isWebSocket := false
		if strings.ToLower(r.Header.Get("Upgrade")) == "websocket" {
			isWebSocket = true
		}

		if isWebSocket {
			// WEBSOCKET STRATEGY: Hijack and bidirectional copy
			hijacker, ok := w.(http.Hijacker)
			if !ok {
				http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
				return
			}
			clientConn, bufrw, err := hijacker.Hijack()
			if err != nil {
				http.Error(w, err.Error(), http.StatusServiceUnavailable)
				return
			}
			defer clientConn.Close()

			// Manual Request writing to avoid touching Body after Hijack/Panic
			// Request Line
			reqLine := fmt.Sprintf("%s %s %s\r\n", r.Method, r.RequestURI, r.Proto)
			if _, err := io.WriteString(ch, reqLine); err != nil {
				log.Printf("Error writing websocket request line: %v", err)
				return
			}
			// Headers
			if err := r.Header.Write(ch); err != nil {
				log.Printf("Error writing websocket headers: %v", err)
				return
			}
			// End of headers
			if _, err := io.WriteString(ch, "\r\n"); err != nil {
				log.Printf("Error writing websocket header terminator: %v", err)
				return
			}

			var wg sync.WaitGroup
			wg.Add(2)

			// Copy existing buffer from hijack + future reads -> backend
			go func() {
				defer func() {
					if r := recover(); r != nil {
						log.Printf("Recovered from panic in WS writer: %v", r)
					}
					wg.Done()
				}()
				if bufrw.Reader.Buffered() > 0 {
					io.CopyN(ch, bufrw, int64(bufrw.Reader.Buffered()))
				}
				io.Copy(ch, clientConn)
			}()

			// Backend -> Browser
			go func() {
				defer func() {
					if r := recover(); r != nil {
						log.Printf("Recovered from panic in WS reader: %v", r)
					}
					wg.Done()
				}()
				io.Copy(clientConn, ch)
			}()

			wg.Wait()
			return
		}

		// STANDARD HTTP STRATEGY: Request/Response Tunneling
		// 1. Write Request to Channel (simulating wire)
		if err := r.Write(ch); err != nil {
			log.Printf("Error writing request to backend: %v", err)
			http.Error(w, "Gateway Error", http.StatusBadGateway)
			return
		}

		// 2. Read Response from Channel (parsing wire)
		resp, err := http.ReadResponse(bufio.NewReader(ch), r)
		if err != nil {
			log.Printf("Error reading response from backend: %v", err)
			http.Error(w, "Bad Gateway", http.StatusBadGateway)
			return
		}
		defer resp.Body.Close()

		// 3. Copy Headers to `w`
		for k, vv := range resp.Header {
			for _, v := range vv {
				w.Header().Add(k, v)
			}
		}
		w.WriteHeader(resp.StatusCode)

		// 4. Copy Body to `w`
		io.Copy(w, resp.Body)

	})

	log.Fatal(http.ListenAndServe(port, nil))
}

func main() {
	rand.Seed(time.Now().UnixNano())

	// 0. Parse Flags
	sshAddr := flag.String("ssh", ":2222", "Address for SSH server to listen on")
	httpAddr := flag.String("http", ":8080", "Address for HTTP proxy to listen on")
	flag.Parse()

	// 1. Auth Token Setup
	authToken := os.Getenv("GROKWAY_TOKEN")
	if authToken == "" {
		authToken = generateSlug(16) // Longer token
		fmt.Printf("\nGenerated Admin Token: %s\n\n", authToken)
		fmt.Println("Use this token to configure your client:")
		fmt.Printf("  grokway client configure --token %s --server <GROKWAY_SERVER_ADDRESS>\n\n", authToken)
	} else {
		fmt.Println("Using configured GROKWAY_TOKEN")
	}

	// 2. SSH Server
	sshServer := &ssh.Server{
		Addr: *sshAddr,
		PasswordHandler: func(ctx ssh.Context, password string) bool {
			return password == authToken
		},
		Handler: func(s ssh.Session) {
			// This handler handles the "Shell" or "Exec" requests.
			// The port forwarding happens via "tcpip-forward" request global handler.

			// Show TUI-ish banner to the SSH client
			io.WriteString(s, "Welcome to GrokwayV2!\n")
			io.WriteString(s, fmt.Sprintf("You are connected as %s\n", s.User()))

			// Keep session open
			select {}
		},
	}

	// Register global request handler for "tcpip-forward"
	// This is where clients say "Please listen on port X"
	sshServer.RequestHandlers = map[string]ssh.RequestHandler{
		"tcpip-forward": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)

			// Check if client requested a specific slug
			requestedSlugsMu.Lock()
			slug, hasRequested := requestedSlugs[conn]
			if hasRequested {
				delete(requestedSlugs, conn)
			}
			requestedSlugsMu.Unlock()

			if !hasRequested {
				slug = generateSlug(8)
			}

			log.Printf("Client requested forwarding. Assigning slug: %s", slug)

			manager.Register(slug, &Tunnel{
				ID:        slug,
				LocalPort: 80, // We assume client forwards to port 80 locally? No,
				// the client says "forward FROM remote port X".
				// The server opens channel TO client when traffic hits X.
				Conn: conn,
			})

			// Monitor connection for closure to cleanup
			go func() {
				conn.Wait()
				log.Printf("Connection closed for slug %s, unregistering...", slug)
				manager.Unregister(slug)
			}()

			// Reply success
			// Payload: port that was bound (uint32)
			// We can return a successful port, e.g. 0 or 80.
			portBytes := []byte{0, 0, 0, 80}
			return true, portBytes
		},
		"cancel-tcpip-forward": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			return true, nil
		},
		"grokway-request-slug": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			slug := string(req.Payload)
			if !slugPattern.MatchString(slug) {
				log.Printf("Invalid slug format requested: %q", slug)
				return false, []byte("invalid slug format (lowercase alphanumeric and hyphens, 3-32 chars)")
			}
			// Check if slug is already in use
			if _, taken := manager.Get(slug); taken {
				log.Printf("Slug %q already in use", slug)
				return false, []byte("slug already in use")
			}
			conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)
			requestedSlugsMu.Lock()
			requestedSlugs[conn] = slug
			requestedSlugsMu.Unlock()
			log.Printf("Client requested slug: %s", slug)
			return true, nil
		},
		"grokway-whoami": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)
			slug, ok := manager.FindSlugByConn(conn)
			if !ok {
				// Not found? Maybe not forwarded yet?
				log.Printf("Whoami requested but no tunnel found for conn %v", conn.RemoteAddr())
				return false, nil
			}
			return true, []byte(slug)
		},
	}

	// Load Host Key (optional for dev, helpful for no warnings)
	// sshServer.SetOption(ssh.HostKeyFile("host.key"))

	fmt.Printf("Starting GrokwayV2 Server on %s\n", *sshAddr)
	fmt.Printf("Starting HTTP Proxy on %s\n", *httpAddr)

	go startHttpProxy(*httpAddr)

	log.Fatal(sshServer.ListenAndServe())
}