package browsergateway

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/coder/websocket"
	"github.com/fosrl/newt/logger"
	"golang.org/x/crypto/ssh"
)

// sshClientMsg is a JSON message sent from the browser to the proxy.
type sshClientMsg struct {
	// type: "auth" | "data" | "resize"
	Type        string `json:"type"`
	Password    string `json:"password,omitempty"`    // used when type="auth"
	PrivateKey  string `json:"privateKey,omitempty"`  // used when type="auth"
	Certificate string `json:"certificate,omitempty"` // used when type="auth"
	Data        string `json:"data,omitempty"`        // used when type="data"
	Cols        uint32 `json:"cols,omitempty"`        // used when type="resize"
	Rows        uint32 `json:"rows,omitempty"`        // used when type="resize"
}

// sshServerMsg is a JSON message sent from the proxy back to the browser.
type sshServerMsg struct {
	// type: "data" | "error"
	Type  string `json:"type"`
	Data  string `json:"data,omitempty"`
	Error string `json:"error,omitempty"`
}

// HandleSSH is an http.HandlerFunc for SSH-over-WebSocket connections.
func (g *Gateway) HandleSSH(w http.ResponseWriter, r *http.Request) {
	logger.Debug("SSH connection request from %s", r.RemoteAddr)
	ctx := r.Context()

	token := r.URL.Query().Get("authToken")

	// "mode=native" (default) connects to the local SSH daemon on this host.
	// "mode=proxy" connects to an arbitrary host+port supplied in query params.
	nativeSSH := r.URL.Query().Get("mode") != "proxy"

	// In proxy mode we also need host + username from query params.
	var target, username string
	if !nativeSSH {
		host := r.URL.Query().Get("host")
		port := r.URL.Query().Get("port")
		username = r.URL.Query().Get("username")
		if host == "" || username == "" {
			http.Error(w, "missing host or username", http.StatusBadRequest)
			return
		}
		if port == "" {
			port = "22"
		}
		sshPort, _ := strconv.Atoi(port)
		if !g.isAllowed("ssh", host, sshPort, token) {
			http.Error(w, "destination not allowed or auth token mismatch", http.StatusForbidden)
			return
		}
		target = net.JoinHostPort(host, port)
	} else {
		// Native SSH mode: validate the token against any registered ssh target.
		if !g.isTokenValid("ssh", token) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		username = r.URL.Query().Get("username")
		if username == "" {
			http.Error(w, "missing username", http.StatusBadRequest)
			return
		}
	}

	ws, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		InsecureSkipVerify: true,
		Subprotocols:       []string{"ssh"},
	})
	if err != nil {
		logger.Debug("SSH websocket upgrade failed: %v", err)
		return
	}
	ws.SetReadLimit(-1)
	defer ws.CloseNow() //nolint:errcheck

	if nativeSSH {
		if err := serveNativeSSHSession(ctx, ws, username, g.sshCreds); err != nil {
			logger.Debug("SSH native session error: %v", err)
		}
	} else {
		if err := serveSSHSession(ctx, ws, target, username, g.authToken); err != nil {
			logger.Debug("SSH session error: %v", err)
		}
	}
}

func serveSSHSession(ctx context.Context, ws *websocket.Conn, target, username, _ string) error {
	// -- Wait for the auth message from the client to get the password --
	_, authBytes, err := ws.Read(ctx)
	if err != nil {
		return fmt.Errorf("read auth message: %w", err)
	}
	var authMsg sshClientMsg
	if err := json.Unmarshal(authBytes, &authMsg); err != nil || authMsg.Type != "auth" {
		return fmt.Errorf("expected auth message, got: %s", authBytes)
	}
	password := authMsg.Password
	privateKey := authMsg.PrivateKey
	certificate := authMsg.Certificate

	// Build the list of auth methods. Certificate+private key takes priority
	// when both are provided, then private key, then password.
	var authMethods []ssh.AuthMethod
	if privateKey != "" {
		signer, err := ssh.ParsePrivateKey([]byte(privateKey))
		if err != nil {
			sendSSHError(ctx, ws, fmt.Sprintf("Failed to parse private key: %v", err))
			return fmt.Errorf("parse private key: %w", err)
		}

		if certificate != "" {
			certPubKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(certificate))
			if err != nil {
				sendSSHError(ctx, ws, fmt.Sprintf("Failed to parse certificate: %v", err))
				return fmt.Errorf("parse certificate: %w", err)
			}

			cert, ok := certPubKey.(*ssh.Certificate)
			if !ok {
				sendSSHError(ctx, ws, "Provided certificate is not a valid SSH certificate")
				return fmt.Errorf("provided certificate is not ssh certificate")
			}

			certSigner, err := ssh.NewCertSigner(cert, signer)
			if err != nil {
				sendSSHError(ctx, ws, fmt.Sprintf("Failed to combine private key and certificate: %v", err))
				return fmt.Errorf("new cert signer: %w", err)
			}
			authMethods = append(authMethods, ssh.PublicKeys(certSigner))
		} else {
			authMethods = append(authMethods, ssh.PublicKeys(signer))
		}
	}
	if password != "" {
		authMethods = append(authMethods, ssh.Password(password))
	}
	if len(authMethods) == 0 {
		sendSSHError(ctx, ws, "No authentication credentials provided")
		return fmt.Errorf("no auth credentials")
	}
	logger.Debug("SSH: connecting to %s as %s", target, username)
	sshCfg := &ssh.ClientConfig{
		User: username,
		Auth: authMethods,
		// HostKeyCallback is intentionally InsecureIgnoreHostKey for this dev
		// proxy. In production, verify against a known-hosts store.
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec
		Timeout:         15 * time.Second,
	}

	sshClient, err := ssh.Dial("tcp", target, sshCfg)
	if err != nil {
		sendSSHError(ctx, ws, fmt.Sprintf("SSH dial failed: %v", err))
		return fmt.Errorf("ssh dial %s: %w", target, err)
	}
	defer sshClient.Close()

	// -- Open an interactive session --
	sess, err := sshClient.NewSession()
	if err != nil {
		sendSSHError(ctx, ws, fmt.Sprintf("Failed to open SSH session: %v", err))
		return fmt.Errorf("ssh new session: %w", err)
	}
	defer sess.Close()

	// Request a PTY.
	if err := sess.RequestPty("xterm-256color", 24, 80, ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 38400,
		ssh.TTY_OP_OSPEED: 38400,
	}); err != nil {
		sendSSHError(ctx, ws, fmt.Sprintf("Failed to request PTY: %v", err))
		return fmt.Errorf("ssh request pty: %w", err)
	}

	stdinPipe, err := sess.StdinPipe()
	if err != nil {
		return fmt.Errorf("ssh stdin pipe: %w", err)
	}
	stdoutPipe, err := sess.StdoutPipe()
	if err != nil {
		return fmt.Errorf("ssh stdout pipe: %w", err)
	}
	stderrPipe, err := sess.StderrPipe()
	if err != nil {
		return fmt.Errorf("ssh stderr pipe: %w", err)
	}

	if err := sess.Shell(); err != nil {
		sendSSHError(ctx, ws, fmt.Sprintf("Failed to start shell: %v", err))
		return fmt.Errorf("ssh shell: %w", err)
	}

	logger.Debug("SSH: session established with %s", target)

	// -- Pump SSH stdout/stderr → WebSocket --
	sessCtx, cancelSess := context.WithCancel(ctx)
	defer cancelSess()

	go func() {
		buf := make([]byte, 4096)
		for {
			n, readErr := stdoutPipe.Read(buf)
			if n > 0 {
				msg := sshServerMsg{Type: "data", Data: string(buf[:n])}
				b, _ := json.Marshal(msg)
				if writeErr := ws.Write(sessCtx, websocket.MessageText, b); writeErr != nil {
					return
				}
			}
			if readErr != nil {
				cancelSess()
				return
			}
		}
	}()

	go func() {
		buf := make([]byte, 4096)
		for {
			n, readErr := stderrPipe.Read(buf)
			if n > 0 {
				msg := sshServerMsg{Type: "data", Data: string(buf[:n])}
				b, _ := json.Marshal(msg)
				if writeErr := ws.Write(sessCtx, websocket.MessageText, b); writeErr != nil {
					return
				}
			}
			if readErr != nil {
				return
			}
		}
	}()

	// -- Pump WebSocket input → SSH stdin / resize --
	for {
		_, msgBytes, readErr := ws.Read(sessCtx)
		if readErr != nil {
			break
		}

		var msg sshClientMsg
		if err := json.Unmarshal(msgBytes, &msg); err != nil {
			continue
		}

		switch msg.Type {
		case "data":
			if _, err := stdinPipe.Write([]byte(msg.Data)); err != nil {
				return fmt.Errorf("write ssh stdin: %w", err)
			}
		case "resize":
			if msg.Cols > 0 && msg.Rows > 0 {
				_ = sess.WindowChange(int(msg.Rows), int(msg.Cols))
			}
		}
	}

	return nil
}

// sendSSHError sends an error message to the browser and logs it.
func sendSSHError(ctx context.Context, ws *websocket.Conn, msg string) {
	logger.Debug("SSH error: %s", msg)
	b, _ := json.Marshal(sshServerMsg{Type: "error", Error: msg})
	_ = ws.Write(ctx, websocket.MessageText, b)
}