package handlers

import (
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh"
	"gorm.io/gorm"

	"ezops/internal/models"
)

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

// WSMessage WebSocket消息结构 // Logan.Li
type WSMessage struct {
	Status int    `json:"status"` // 0=data, 1=resize
	Data   string `json:"data"`
	Cols   int    `json:"cols"`
	Rows   int    `json:"rows"`
}

// WebShellHandler WebShell处理器
type WebShellHandler struct {
	db      *gorm.DB
	clients sync.Map
}

// NewWebShellHandler 创建WebShell处理器
func NewWebShellHandler(db *gorm.DB) *WebShellHandler {
	return &WebShellHandler{db: db}
}

// formatSSHError 将 Go SSH 错误转换为友好的中文提示 // Logan.Li
func formatSSHError(err error) string {
	errStr := err.Error()
	switch {
	case contains(errStr, "i/o timeout"):
		return "连接超时（目标主机可能在防火墙后或网络不通）"
	case contains(errStr, "connection refused"):
		return "连接被拒绝（SSH 服务未运行或端口错误）"
	case contains(errStr, "no route to host"):
		return "无法路由到目标主机（IP 地址不可达）"
	case contains(errStr, "authentication failed") || contains(errStr, "unable to authenticate"):
		return "认证失败（用户名或密码错误）"
	case contains(errStr, "handshake"):
		return "SSH 握手失败（目标可能不是 SSH 服务）"
	case contains(errStr, "dns") || contains(errStr, "no such host"):
		return "域名解析失败（IP 地址无效）"
	case contains(errStr, "ssh: must specify"):
		return "认证方式错误（未配置密码或密钥）"
	default:
		return errStr
	}
}

func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
		(len(s) > 0 && len(substr) > 0 && indexOf(s, substr) >= 0))
}

func indexOf(s, substr string) int {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return i
		}
	}
	return -1
}

// Terminal WebSocket终端连接 — 真实SSH，带终端会话录制 // Logan.Li
func (h *WebShellHandler) Terminal(c *gin.Context) {
	serverID := c.Query("id")
	if serverID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "missing server id"})
		return
	}

	var server models.Server
	if err := h.db.First(&server, serverID).Error; err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "server not found"})
		return
	}

	// 先升级 WebSocket — 这样前端能看到连接过程
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("[WebShell] WebSocket升级失败: %v", err)
		return
	}
	defer conn.Close()

	// 创建终端会话记录（终端审计） // Logan.Li
	var userIDPtr *uint
	if uid, exists := c.Get("userID"); exists {
		if uidVal, ok := uid.(uint); ok {
			userIDPtr = &uidVal
		}
	}
	username, _ := c.Get("username")
	usernameStr, _ := username.(string)

	termSession := models.TerminalSession{
		UserID:     userIDPtr,
		Username:   usernameStr,
		ServerID:   server.ID,
		ServerName: server.Hostname,
		ServerIP:   server.IP,
		ServerPort: server.Port,
		ClientIP:   c.ClientIP(),
		Status:     "active",
		StartTime:  time.Now(),
	}
	if err := h.db.Create(&termSession).Error; err != nil {
		log.Printf("[WebShell] 创建会话记录失败: %v", err)
	}

	// 终端录制缓冲（限制 1MB 避免内存问题） // Logan.Li
	const maxRecordSize = 1024 * 1024
	var recordBuf strings.Builder

	// finalizeSession 结束会话记录（统一收尾）
	finalizeSession := func(status string) {
		endTime := time.Now()
		duration := int(endTime.Sub(termSession.StartTime).Seconds())
		h.db.Model(&termSession).Updates(map[string]interface{}{
			"status":      status,
			"duration":    duration,
			"end_time":    &endTime,
			"record_data": recordBuf.String(),
			"record_size": recordBuf.Len(),
		})
	}

	conn.WriteJSON(WSMessage{
		Status: 0,
		Data:   "\r\n\x1b[36m正在连接 " + server.Hostname + " (" + server.IP + ":" + strconv.Itoa(server.Port) + ")...\x1b[0m\r\n",
	})

	// SSH HostKey 验证 — Trust-On-First-Use (TOFU).
	// 首次连接记录 fingerprint 到日志；后续连接通过已知 HostKey 校验。
	knownHosts := make(map[string]ssh.PublicKey)
	var knownMu sync.Mutex
	sshConfig := &ssh.ClientConfig{
		User:    server.Username,
		Timeout: 10 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			fp := ssh.FingerprintSHA256(key)
			knownMu.Lock()
			defer knownMu.Unlock()
			if existing, ok := knownHosts[hostname]; ok {
				if ssh.FingerprintSHA256(existing) == fp {
					return nil
				}
				return fmt.Errorf("SSH host key mismatch for %s: expected %s, got %s",
					hostname, ssh.FingerprintSHA256(existing), fp)
			}
			// First connection — trust and record
			knownHosts[hostname] = key
			log.Printf("[WebShell] TOFU: trusting new host key for %s (%s): %s",
				server.Hostname, hostname, fp)
			return nil
		},
	}

	switch server.AuthType {
	case "key":
		signer, err := ssh.ParsePrivateKey([]byte(server.SSHKey))
		if err != nil {
			sshConfig.Auth = []ssh.AuthMethod{ssh.Password(server.Password)}
		} else {
			sshConfig.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
		}
	default:
		sshConfig.Auth = []ssh.AuthMethod{
			ssh.Password(server.Password),
		}
	}

	// 连接 SSH（带10秒超时）
	addr := server.IP + ":" + strconv.Itoa(server.Port)
	sshClient, err := ssh.Dial("tcp", addr, sshConfig)
	if err != nil {
		log.Printf("[WebShell] SSH连接失败 %s: %v", addr, err)
		friendlyErr := formatSSHError(err)
		errMsg := "\r\n\x1b[31m连接失败: " + friendlyErr + "\x1b[0m\r\n" +
			"\x1b[33m──────────────────────────\x1b[0m\r\n" +
			"\x1b[33m目标:\x1b[0m " + server.Hostname + " (" + server.IP + ":" + strconv.Itoa(server.Port) + ")\r\n" +
			"\x1b[33m账号:\x1b[0m " + server.Username + "\r\n" +
			"\x1b[33m排查:\x1b[0m ① 检查 IP/端口是否正确  ② 确认 SSH 服务运行中  ③ 验证账号密码\r\n" +
			"\x1b[33m──────────────────────────\x1b[0m\r\n"
		conn.WriteJSON(WSMessage{Status: 0, Data: errMsg})
		recordBuf.WriteString(errMsg)
		finalizeSession("failed")
		conn.Close()
		return
	}
	defer sshClient.Close()

	log.Printf("[WebShell] SSH已连接 %s (%s)", server.Hostname, addr)

	// 创建 SSH session
	session, err := sshClient.NewSession()
	if err != nil {
		log.Printf("[WebShell] 创建session失败: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31m创建SSH会话失败\x1b[0m\r\n"})
		finalizeSession("failed")
		return
	}
	defer session.Close()

	// 请求 PTY
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}
	if err := session.RequestPty("xterm-256color", 24, 80, modes); err != nil {
		log.Printf("[WebShell] PTY请求失败: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mPTY请求失败\x1b[0m\r\n"})
		finalizeSession("failed")
		return
	}

	// 获取管道
	stdinPipe, err := session.StdinPipe()
	if err != nil {
		finalizeSession("failed")
		return
	}
	stdoutPipe, err := session.StdoutPipe()
	if err != nil {
		finalizeSession("failed")
		return
	}

	// 启动 shell
	if err := session.Shell(); err != nil {
		log.Printf("[WebShell] Shell启动失败: %v", err)
		finalizeSession("failed")
		return
	}

	// 发送欢迎消息
	welcome := "\r\n\x1b[32mAIOps WebShell\x1b[0m — 已连接 " + server.Hostname + " (" + server.IP + ")\r\n\r\n"
	conn.WriteJSON(WSMessage{
		Status: 0,
		Data:   welcome,
	})
	recordBuf.WriteString(welcome)

	done := make(chan struct{})

	// goroutine: SSH stdout -> WebSocket（同时录制输出） // Logan.Li
	go func() {
		defer close(done)
		buf := make([]byte, 8192)
		for {
			n, err := stdoutPipe.Read(buf)
			if err != nil {
				discMsg := "\r\n\x1b[31m[远程连接已断开]\x1b[0m\r\n"
				conn.WriteJSON(WSMessage{Status: 0, Data: discMsg})
				recordBuf.WriteString(discMsg)
				break
			}
			if n > 0 {
				conn.WriteJSON(WSMessage{Status: 0, Data: string(buf[:n])})
				// 录制终端输出（受 1MB 上限保护）
				if recordBuf.Len() < maxRecordSize {
					remain := maxRecordSize - recordBuf.Len()
					if n > remain {
						recordBuf.Write(buf[:remain])
					} else {
						recordBuf.Write(buf[:n])
					}
				}
			}
		}
	}()

	// 主循环: WebSocket -> SSH stdin
	for {
		msgType, data, err := conn.ReadMessage()
		if err != nil {
			break
		}

		if msgType == websocket.TextMessage {
			var msg WSMessage
			if err := json.Unmarshal(data, &msg); err != nil {
				continue
			}

			switch msg.Status {
			case 0: // 数据输入
				stdinPipe.Write([]byte(msg.Data))
			case 1: // resize
				if msg.Cols > 0 && msg.Rows > 0 {
					session.WindowChange(msg.Rows, msg.Cols)
				}
			}
		}
	}

	// 等待 stdout goroutine 结束
	<-done
	stdinPipe.Close()
	finalizeSession("closed")
	log.Printf("[WebShell] 断开 %s (%s)", server.Hostname, addr)
}