package main

import (
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"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),
}

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

	http.Handle("/assets/", http.StripPrefix("/assets/", 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]

		tunnel, ok := manager.Get(slug)
		if !ok {
			// Serve 404 page
			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.

		// Manual channel opening:
		if tunnel.Conn == nil {
			http.Error(w, "Tunnel Broken", http.StatusBadGateway)
			return
		}

		payload := make([]byte, 0) // dynamic payload constructing
		// Open channel to client
		// "forwarded-tcpip" arguments:
		// string    address that was connected
		// uint32    port that was connected
		// string    originator IP address
		// uint32    originator port

		// For simplicity, we just say we are connecting to 127.0.0.1:LocalPort
		// And originating from 127.0.0.1:0

		destHost := "0.0.0.0"
		destPort := tunnel.LocalPort
		srcHost := "127.0.0.1"
		srcPort := uint32(12345) // Random source port, cannot be 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)
			http.Error(w, "Failed to connect to client", http.StatusBadGateway)
			return
		}
		defer ch.Close()
		go gossh.DiscardRequests(reqs)

		// Pipe HTTP request to the channel?
		// No, we can't just pipe the raw HTTP structure unless the client expects it.
		// If the client is just a TCP forwarder (ssh -R), it expects a raw TCP stream.
		// So we need to act as a TCP proxy.
		// But we are in an http.Handler.
		// We can hijack the connection or just Re-serialize the HTTP request.

		// Re-serializing is safer for now.
		// Or we use Hijack to get the raw TCP conn.

		// Let's iterate: just write the Request to the channel
		// The client will see it as raw bytes on the socket.
		// This works if the client is forwarding to a web server.

		err = r.Write(ch) // Writes the request wire representation
		if err != nil {
			http.Error(w, "Failed to write request: "+err.Error(), http.StatusInternalServerError)
			return
		}

		// Read response from channel and write to w
		// This is tricky because we need to parse the response to set headers correctly in w.
		// Or we can just use Hijack if we want to be a transparent TCP proxy.
		// But Hijack breaks HTTP/2 and some middleware.
		// Let's use `http.ReadResponse`

		// Problem: `r.Write(ch)` might not close the write side, so the server might wait.
		// but HTTP 1.1 usually has content-length.

		// Using bufio to read the response from the channel.
		// resp, err := http.ReadResponse(bufio.NewReader(ch), r)
		// ... copy headers, body ...

		// For MVP, hijacking is often easiest for "Tunneling"
		hijacker, ok := w.(http.Hijacker)
		if !ok {
			http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
			return
		}
		clientConn, _, err := hijacker.Hijack()
		if err != nil {
			http.Error(w, err.Error(), http.StatusServiceUnavailable)
			return
		}
		defer clientConn.Close()

		// bidirectional copy
		var wg sync.WaitGroup
		wg.Add(2)

		// We already consumed `r` somewhat? No, Hijack takes over.
		// But `http.HandleFunc` might have read the headers.
		// The `clientConn` is the raw TCP connection from the Browser/ExternalUser to OUR Proxy.
		// We need to replay the request headers!
		// Because `http.Server` already read them.

		// Actually, `r.Write(ch)` is correct for sending the request to the backend.
		// But then we need to pipe `ch` output back to `w`.
		// And since we used Hijack, we are responsible for the entire response.

		// So:
		// 1. Write `r` to `ch`.
		// 2. IO Copy `ch` -> `clientConn`.

		go func() {
			defer wg.Done()
			r.Write(ch) // Send the request to valid backend
			// Also forward any extra body if not fully read?
			// r.Write handles body.
		}()

		go func() {
			defer wg.Done()
			io.Copy(clientConn, ch) // Copy response from backend to user
		}()

		// Wait? No, if we wait for r.Write, it might block if the body is large.
		// Actually, standard proxying is complex.
		// Simplest "Grokway" logic:
		// User runs: ssh -R 80:localhost:8080
		// Server listens on a random port (or we map virtual host).

		// Let's ignore Hijack for a moment and try standard Request/Response.
		// Implementation detail for later.

		// Connection is hijacked, do not write to w

		wg.Wait()

	})

	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) {
			// Parse payload
			// string address to bind (usually empty or 0.0.0.0 or 127.0.0.1)
			// uint32 port number to bind

			// For Grokway, we ignore the requested port and assign a random subdomain/slug
			// Or we use the requested port if valid?
			// Let's assume we ignore it and generate a slug,
			// OR we use the port as the "slug" if it's special.

			// But wait, the client needs to know what we assigned!
			// Standard SSH response to tcpip-forward contains the BOUND PORT.
			// If we return a port, the client knows.
			// If we want to return a URL, standard SSH doesn't have a field for that.
			// But we can write it to the Session (stdout).

			// Let's accept any request and map it to a random slug.
			slug := generateSlug(8)

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

			conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)

			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-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())
}