package web_console

import (
	"crypto/ed25519"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/gorilla/websocket"
	"github.com/pingcap/log"
	"go.uber.org/zap"
	"golang.org/x/crypto/ssh"
)

const (
	// Time allowed to write or read a message.
	messageWait = 10 * time.Second

	// WebSocket buffer size. This is not a hard message-size limit.
	maxMessageSize = 512

	// sshStreamBufferSize is the buffer size used when reading SSH stdout/stderr before sending to the WebSocket peer.
	sshStreamBufferSize = 32 * 1024

	// maxWSReadLimit is the maximum size in bytes for a WebSocket message read from the peer.
	// Control messages (e.g. resize) are expected to be tiny; keep this reasonably small.
	maxWSReadLimit = 64 * 1024

	// Terminal size sanity limits to reject obviously invalid/malicious resize values.
	maxTerminalRows = 4096
	maxTerminalCols = 8192
)

var terminalModes = ssh.TerminalModes{
	ssh.ECHO:          1,     // enable echoing (different from the example in docs)
	ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
	ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  maxMessageSize,
	WriteBufferSize: maxMessageSize,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// activeWebConsoleSessions tracks the number of currently active WebSocket-based SSH sessions.
var activeWebConsoleSessions int64

// getActiveWebConsoleSessions returns the current number of active sessions.
func getActiveWebConsoleSessions() int64 {
	return atomic.LoadInt64(&activeWebConsoleSessions)
}

type windowSize struct {
	High  int `json:"high"`
	Width int `json:"width"`
}

type sshClient struct {
	conn       *websocket.Conn
	addr       string
	user       string
	keyfile    string
	client     *ssh.Client
	sess       *ssh.Session
	sessIn     io.WriteCloser
	sessOut    io.Reader
	sessErr    io.Reader
	writeMu    sync.Mutex
	closeSig   chan struct{}
	outputDone int32
}

func (c *sshClient) signalClose() {
	if c.closeSig == nil {
		return
	}
	select {
	case c.closeSig <- struct{}{}:
	default:
	}
}

func (c *sshClient) signalOutputDone() {
	// Close the bridge only after both stdout and stderr streams finish cleanly,
	// so we don't drop buffered data from the other stream.
	if atomic.AddInt32(&c.outputDone, 1) == 2 {
		c.signalClose()
	}
}

func (c *sshClient) writeBinaryMessage(payload []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	c.conn.SetWriteDeadline(time.Now().Add(messageWait))
	return c.conn.WriteMessage(websocket.BinaryMessage, payload)
}

func (c *sshClient) getWindowSize() (wdSize *windowSize, err error) {
	c.conn.SetReadDeadline(time.Now().Add(messageWait))

	wdSize = new(windowSize)
	// Wait for the first window-size control message ("1:"+JSON).
	for {
		msgType, connReader, err := c.conn.NextReader()
		if err != nil {
			return nil, fmt.Errorf("conn.NextReader: %w", err)
		}
		if msgType != websocket.TextMessage {
			if err := drainWSMessage(connReader); err != nil {
				return nil, fmt.Errorf("drain non-text message: %w", err)
			}
			continue
		}

		ok, err := parseWindowSizeControlMessage(connReader, wdSize)
		if err != nil {
			return nil, err
		}
		if ok {
			return wdSize, nil
		}
	}
}

func drainWSMessage(r io.Reader) error {
	var buf [256]byte
	for {
		_, err := r.Read(buf[:])
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}
}

func parseWindowSizeControlMessage(r io.Reader, dst *windowSize) (bool, error) {
	var prefix [2]byte
	if _, err := io.ReadFull(r, prefix[:]); err != nil {
		if drainErr := drainWSMessage(r); drainErr != nil {
			return false, fmt.Errorf("drain short text message: %w", drainErr)
		}
		return false, nil
	}
	if prefix[0] != '1' || prefix[1] != ':' {
		if err := drainWSMessage(r); err != nil {
			return false, fmt.Errorf("drain non-control text message: %w", err)
		}
		return false, nil
	}

	dec := json.NewDecoder(r)
	if err := dec.Decode(dst); err != nil {
		return false, fmt.Errorf("json.Decode: %w", err)
	}
	if err := dec.Decode(&struct{}{}); err != io.EOF {
		if err == nil {
			return false, fmt.Errorf("window size message has extra JSON value")
		}
		return false, fmt.Errorf("window size message has trailing data: %w", err)
	}
	if err := validateWindowSize(*dst); err != nil {
		return false, err
	}
	return true, nil
}

func validateWindowSize(w windowSize) error {
	if w.High <= 0 || w.Width <= 0 {
		return fmt.Errorf("invalid window size: high=%d width=%d", w.High, w.Width)
	}
	if w.High > maxTerminalRows || w.Width > maxTerminalCols {
		return fmt.Errorf("window size out of range: high=%d width=%d", w.High, w.Width)
	}
	return nil
}

func (c *sshClient) wsWrite() (retErr error) {
	defer func() {
		if retErr != nil {
			c.signalClose()
			return
		}
		c.signalOutputDone()
	}()

	data := make([]byte, sshStreamBufferSize)

	for {
		n, readErr := c.sessOut.Read(data)
		if n > 0 {
			if err := c.writeBinaryMessage(data[:n]); err != nil {
				return fmt.Errorf("conn.WriteMessage: %w", err)
			}
		}
		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				return nil
			}
			return fmt.Errorf("sessOut.Read: %w", readErr)
		}
	}
}

func (c *sshClient) wsWriteStderr() (retErr error) {
	defer func() {
		if retErr != nil {
			c.signalClose()
			return
		}
		c.signalOutputDone()
	}()

	data := make([]byte, sshStreamBufferSize)

	for {
		n, readErr := c.sessErr.Read(data)
		if n > 0 {
			if err := c.writeBinaryMessage(data[:n]); err != nil {
				return fmt.Errorf("conn.WriteMessage: %w", err)
			}
		}
		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				return nil
			}
			return fmt.Errorf("sessErr.Read: %w", readErr)
		}
	}
}

func (c *sshClient) wsRead() error {
	defer func() {
		c.signalClose()
	}()

	var zeroTime time.Time
	c.conn.SetReadDeadline(zeroTime)

	for {
		msgType, connReader, err := c.conn.NextReader()
		if err != nil {
			return fmt.Errorf("conn.NextReader: %w", err)
		}
		if msgType == websocket.BinaryMessage {
			if _, err := io.Copy(c.sessIn, connReader); err != nil {
				return fmt.Errorf("io.Copy: %w", err)
			}
			continue
		}
		if msgType != websocket.TextMessage {
			if err := drainWSMessage(connReader); err != nil {
				return fmt.Errorf("drain unknown message: %w", err)
			}
			continue
		}

		var wdSize windowSize
		ok, err := parseWindowSizeControlMessage(connReader, &wdSize)
		if err != nil {
			return err
		}
		if !ok {
			continue
		}
		if err := c.sess.WindowChange(wdSize.High, wdSize.Width); err != nil {
			return fmt.Errorf("sess.WindowChange: %w", err)
		}
	}
}

func (c *sshClient) bridgeWSAndSSH() {
	defer c.conn.Close()

	wdSize, err := c.getWindowSize()
	if err != nil {
		log.Info("bridgeWSAndSSH: getWindowSize:", zap.Error(err))
		return
	}

	key, err := os.ReadFile(c.keyfile)
	if err != nil {
		log.Info("bridgeWSAndSSH: os.ReadFile:", zap.Error(err))
		return
	}
	privateKey, err := ssh.ParsePrivateKey(key)
	if err != nil {
		log.Info("bridgeWSAndSSH: ssh.ParsePrivateKey:", zap.Error(err))
		return
	}
	auth := ssh.PublicKeys(privateKey)

	config := &ssh.ClientConfig{
		User:            c.user,
		Auth:            []ssh.AuthMethod{auth},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         time.Second * 5,
	}
	c.client, err = ssh.Dial("tcp", c.addr, config)
	if err != nil {
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to dial: "+err.Error()))
		c.conn.Close()
		return
	}
	defer c.client.Close()

	c.sess, err = c.client.NewSession()
	if err != nil {
		log.Error("failed to create ssh session", zap.Error(err))
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to create ssh session: "+err.Error()))
		c.conn.Close()
		return
	}
	defer c.sess.Close()

	c.sessOut, err = c.sess.StdoutPipe()
	if err != nil {
		log.Error("failed to create stdout pipe", zap.Error(err))
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to create stdout pipe: "+err.Error()))
		c.conn.Close()
		return
	}

	c.sessErr, err = c.sess.StderrPipe()
	if err != nil {
		log.Error("failed to create stderr pipe", zap.Error(err))
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to create stderr pipe: "+err.Error()))
		c.conn.Close()
		return
	}

	c.sessIn, err = c.sess.StdinPipe()
	if err != nil {
		log.Error("failed to create stdin pipe", zap.Error(err))
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to create stdin pipe: "+err.Error()))
		c.conn.Close()
		return
	}
	defer c.sessIn.Close()

	if err := c.sess.RequestPty("xterm", wdSize.High, wdSize.Width, terminalModes); err != nil {
		log.Error("failed to request pty", zap.Error(err))
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to request pty: "+err.Error()))
		c.conn.Close()
		return
	}

	if err := c.sess.Shell(); err != nil {
		log.Error("failed to start shell", zap.Error(err))
		c.conn.WriteMessage(websocket.TextMessage, []byte("failed to start shell: "+err.Error()))
		c.conn.Close()
		return
	}

	log.Info("started a login shell on the remote host")
	defer log.Info("closed a login shell on the remote host")

	go func() {
		if err := c.wsRead(); err != nil {
			log.Info("bridgeWSAndSSH: wsRead:", zap.Error(err))
		}
	}()

	go func() {
		if err := c.wsWrite(); err != nil {
			log.Info("bridgeWSAndSSH: wsWrite:", zap.Error(err))
		}
	}()

	go func() {
		if err := c.wsWriteStderr(); err != nil {
			log.Info("bridgeWSAndSSH: wsWriteStderr:", zap.Error(err))
		}
	}()

	<-c.closeSig
}

type sshHandler struct {
	jwtPublicKey     ed25519.PublicKey
	sshIdentityFiles map[string]string
}

// webSocket handles WebSocket requests for SSH from the clients.
func (h *sshHandler) webSocket(w http.ResponseWriter, req *http.Request) {
	remoteSshConfig, err := h.verifyJwt(req)
	if err != nil {
		log.Info("failed to verify jwt", zap.Error(err))
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	sshIdentityFile, err := h.getSshIdentityFile(remoteSshConfig.Region)
	if err != nil {
		log.Error("failed to get ssh identity file", zap.Error(err))
		return
	}

	conn, err := upgrader.Upgrade(w, req, nil)
	if err != nil {
		log.Error("failed to upgrade to websocket", zap.Error(err))
		return
	}
	conn.SetReadLimit(maxWSReadLimit)

	sshPort := remoteSshConfig.MachinePort
	if sshPort == 0 {
		sshPort = 22
	}
	sshCli := &sshClient{
		conn:     conn,
		addr:     net.JoinHostPort(remoteSshConfig.MachineIp, strconv.Itoa(sshPort)),
		user:     remoteSshConfig.MachineUser,
		keyfile:  sshIdentityFile,
		closeSig: make(chan struct{}, 1),
	}

	// Increase active session counter on new connection.
	atomic.AddInt64(&activeWebConsoleSessions, 1)
	go func() {
		defer atomic.AddInt64(&activeWebConsoleSessions, -1)
		sshCli.bridgeWSAndSSH()
	}()
}

type RemoteSshConfig struct {
	Id          int    `json:"id"`
	UserId      int    `json:"user_id"`
	Region      string `json:"region"`
	MachineIp   string `json:"m_ip"`
	MachineUser string `json:"m_user"`
	MachinePort int    `json:"m_port"`
	ExpiredAt   int    `json:"exp"`
}

func (h *sshHandler) verifyJwt(req *http.Request) (*RemoteSshConfig, error) {
	token := req.URL.Query().Get("auth")

	parsedToken, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		if token.Method.Alg() != jwt.SigningMethodEdDSA.Alg() {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Method.Alg())
		}
		return h.jwtPublicKey, nil
	})
	if err != nil {
		return nil, fmt.Errorf("jwt.Parse: %w", err)
	}

	if !parsedToken.Valid {
		return nil, fmt.Errorf("token is invalid")
	}

	if claims, ok := parsedToken.Claims.(jwt.MapClaims); ok && parsedToken.Valid {
		data, err := json.Marshal(claims)
		if err != nil {
			return nil, fmt.Errorf("json.Marshal: %w", err)
		}
		log.Info("claims:", zap.String("claims", string(data)))
		var remoteSshConfig RemoteSshConfig
		if err := json.Unmarshal(data, &remoteSshConfig); err != nil {
			return nil, fmt.Errorf("json.Unmarshal: %w", err)
		}
		now := time.Now().Unix()
		if remoteSshConfig.ExpiredAt < int(now) {
			return nil, fmt.Errorf("token is expired")
		}
		return &remoteSshConfig, nil
	}
	return nil, fmt.Errorf("token is invalid")
}

func (h *sshHandler) getSshIdentityFile(region string) (string, error) {
	sshIdentityFile, ok := h.sshIdentityFiles[region]
	if !ok {
		return "", fmt.Errorf("ssh identity file not found for region: %s", region)
	}
	return sshIdentityFile, nil
}