package controller

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"gitee.com/linqwen/momo/app/srv/shell/model"
	"gitee.com/linqwen/momo/app/srv/shell/schema"
	"gitee.com/linqwen/momo/app/srv/shell/service"
	"gitee.com/linqwen/momo/base"
	"gitee.com/linqwen/momo/igin"
	"gitee.com/linqwen/momo/rbac"
	"github.com/gorilla/websocket"
)

type shellcli struct {
	base.BaseController[model.ShellcliEntity, schema.ShellcliCreateDTO, schema.ShellcliQueryDTO, schema.ShellcliVO]
}

func newShellcliService() *shellcli {
	return &shellcli{BaseController: *base.NewBaseController[model.ShellcliEntity, schema.ShellcliCreateDTO, schema.ShellcliQueryDTO, schema.ShellcliVO](service.NewShellcliService())}
}

func ShellcliRouters(router igin.IRouterGroup) {
	c := newShellcliService()
	r := router.Group("shell")
	{
		r.GET("/", c.PageByDTO)
		r.GET("/ws", rbac.Auth.HasPermit("shell"), c.wsHandler)
	}
}

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

func ensureValidUTF8(s string) string {
	if utf8.ValidString(s) {
		return s
	}
	return strings.Map(func(r rune) rune {
		if r == 0 || r == 0xFFFD || r > 0x10FFFF {
			return -1
		}
		return r
	}, s)
}

func formatOutput(output string) string {
	// 替换 Windows 特有的换行符为标准的换行符
	output = strings.ReplaceAll(output, "\r\n", "\n")
	// 确保每行都正确换行
	lines := strings.Split(output, "\n")
	formattedLines := make([]string, len(lines))
	for i, line := range lines {
		// 去除多余的空格和制表符
		formattedLines[i] = strings.TrimSpace(line)
	}
	return strings.Join(formattedLines, "\n")
}

// wsHandler 处理WebSocket请求
func (bc *shellcli) wsHandler(c igin.IContext) {
	log.Println("Handling WebSocket request")
	conn, err := upgrader.Upgrade(c.Writer(), c.Request(), nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}
	defer conn.Close()

	log.Println("WebSocket connection established")

	// 创建一个带超时和取消功能的上下文
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	// 存储当前命令的上下文和互斥锁，以便可以从其他地方安全访问
	commandCtxs := make(map[string]context.CancelFunc)
	commandMutex := &sync.Mutex{}

	// 开始处理WebSocket消息
	for {
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("read error: %v", err)
			}
			return // 如果读取消息出错，则退出循环
		}
		log.Printf("recv: %s, messageType: %v\n", message, messageType)

		// 检查是否是取消命令
		if string(message) == "stop" {
			// 取消当前所有命令
			commandMutex.Lock()
			for id, cancelFunc := range commandCtxs {
				cancelFunc()
				delete(commandCtxs, id)
			}
			commandMutex.Unlock()
			log.Println("All commands cancelled by user")
			conn.WriteMessage(websocket.TextMessage, []byte("All commands were cancelled."))
			continue
		}

		// 生成唯一ID用于标识命令
		cmdID := fmt.Sprintf("%d", time.Now().UnixNano())

		// 为每个命令创建一个新的上下文
		cmdCtx, cmdCancel := context.WithCancel(ctx)

		// 将命令的取消函数添加到映射中
		commandMutex.Lock()
		commandCtxs[cmdID] = cmdCancel
		commandMutex.Unlock()

		// 执行命令
		go handleCommand(cmdCtx, string(message), conn, cmdCancel, cmdID, commandMutex, commandCtxs)
	}
}

// handleCommand 处理单个命令的执行
func handleCommand(ctx context.Context, cmd string, conn *websocket.Conn, cancel context.CancelFunc, cmdID string, commandMutex *sync.Mutex, commandCtxs map[string]context.CancelFunc) {
	var command *exec.Cmd
	var stdout, stderr io.ReadCloser
	var err error

	switch runtime.GOOS {
	case "windows":
		command = exec.CommandContext(ctx, "cmd", "/C", cmd)
	case "linux", "darwin":
		command = exec.CommandContext(ctx, "sh", "-c", cmd)
	default:
		err = fmt.Errorf("unsupported OS: %s", runtime.GOOS)
		conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		return
	}

	// Capture the output streams before starting the command.
	stdout, err = command.StdoutPipe()
	if err != nil {
		log.Printf("error getting StdoutPipe: %s", err)
		conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		return
	}
	stderr, err = command.StderrPipe()
	if err != nil {
		log.Printf("error getting StderrPipe: %s", err)
		conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		return
	}

	// Start the command after capturing the output streams.
	if err = command.Start(); err != nil {
		log.Printf("error starting command: %s", err)
		conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		return
	}

	// 合并标准输出和标准错误流
	multiReader := io.MultiReader(stdout, stderr)

	// 读取命令输出并发送给客户端
	reader := bufio.NewReader(multiReader)
	for {
		select {
		case <-ctx.Done():
			// 如果上下文被取消（超时或收到取消指令），终止命令
			if err = command.Process.Kill(); err != nil {
				log.Printf("error killing command: %s", err)
			}

			// 发送取消消息
			conn.WriteMessage(websocket.TextMessage, []byte("Command was cancelled."))

			// 从命令映射中移除该命令
			commandMutex.Lock()
			delete(commandCtxs, cmdID)
			commandMutex.Unlock()

			// 确保 goroutine 退出
			return
		default:
			line, err := reader.ReadString('\n')
			if err != nil {
				if err == io.EOF {
					// 命令执行完毕
					if err = command.Wait(); err != nil {
						log.Printf("command finished with error: %s", err)
						conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
					} else {
						log.Println("command finished successfully")
					}

					// 从命令映射中移除该命令
					commandMutex.Lock()
					delete(commandCtxs, cmdID)
					commandMutex.Unlock()

					// 确保 goroutine 退出
					return
				}
				log.Printf("error reading output: %s", err)
				return
			}
			// 确保输出是有效的 UTF-8 编码
			validLine := ensureValidUTF8(line)
			// 格式化输出
			formattedLine := formatOutput(validLine)
			conn.WriteMessage(websocket.TextMessage, []byte(formattedLine))
		}
	}
}