package controller

import (
	"context"
	"dashboard/utils"
	"net/http"
	"os/exec"

	"github.com/creack/pty"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"golang.org/x/sync/errgroup"
)

func TerminalHandle() func(c *gin.Context) error {
	upgrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true // 允许跨域
		},
	}
	return func(c *gin.Context) error {
		log, _ := utils.GetLogFromContext(c)

		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			log.Sugar().Errorf("WebSocket upgrade failed: %v", err)
			return err
		}
		defer conn.Close()

		cmd := exec.Command("bash")
		ptyFile, err := pty.Start(cmd)
		if err != nil {
			log.Sugar().Errorf("Failed to start pty: %v", err)
			return err
		}
		defer ptyFile.Close()

		wg, ctx := errgroup.WithContext(context.Background())

		wg.Go(func() error {
			defer log.Sugar().Errorf("task read over")
			for {
				select {
				case <-ctx.Done():
					log.Sugar().Errorf("task canceld read")
					return ctx.Err()
				default:
				}

				_, msg, err := conn.ReadMessage()
				if err != nil {
					log.Sugar().Errorf("WebSocket read error: %v", err)
					return err
				}
				_, err = ptyFile.Write(msg)
				if err != nil {
					log.Sugar().Errorf("PTY write error: %v", err)
					return err
				}
			}
		})

		wg.Go(func() error {
			defer log.Sugar().Errorf("task write over")
			buf := make([]byte, 1024)
			for {
				select {
				case <-ctx.Done():
					log.Sugar().Errorf("task canceld write")
					return ctx.Err()
				default:
				}

				n, err := ptyFile.Read(buf)
				if err != nil {
					log.Sugar().Errorf("PTY read error: %v", err)
					return err
				}
				err = conn.WriteMessage(websocket.BinaryMessage, buf[:n])
				if err != nil {
					log.Sugar().Errorf("WebSocket write error: %v", err)
					return err
				}
			}
		})

		if err := wg.Wait(); err != nil {
			log.Sugar().Errorf("WebSocket error: %v", err)
		}

		return nil
	}
}