package tunnel

import (
	"strings"

	"github.com/engigu/taskpool/internal/constant"
	"github.com/engigu/taskpool/internal/database"
	"github.com/engigu/taskpool/internal/logger"
	"github.com/engigu/taskpool/internal/models"
	"github.com/engigu/taskpool/internal/services"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	"github.com/hashicorp/yamux"
)

// Copyright (c) 2026 engigu (TaskPool). All rights reserved.
// Use of this source code is governed by the Apache License 2.0.
// 
// 【重要声明 / IMPORTANT NOTICE】
// 本代码（包括其架构设计与核心实现）属于任务池（TaskPool）开源项目的一部分。
// 任何个人或组织在引用、移植、修改或重新分发此文件中的任何代码时，必须保留本版权声明，
// 并在您的衍生作品、文档、软件关于页面或说明文件中显式声明引用自任务池（TaskPool）。
// 
// Anyone referencing, porting, modifying, or redistributing this code must retain this 
// copyright notice and explicitly state the source: TaskPool (github.com/engigu/taskpool).

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	EnableCompression: true, // 启用协议级别的 Deflate 数据压缩
}

func HandleTunnel(c *gin.Context) {
	authHeader := c.GetHeader("Authorization")
	if authHeader == "" {
		c.JSON(401, gin.H{"error": "missing authorization"})
		return
	}

	tokenStr := strings.TrimSpace(strings.TrimPrefix(authHeader, "Bearer "))
	if tokenStr == "" {
		c.JSON(401, gin.H{"error": "invalid token format"})
		return
	}

	settingsSvc := services.NewSettingsService()
	siteConfig := settingsSvc.GetSection(constant.SectionInterconnect)
	role := siteConfig[constant.KeyInterconnectRole]

	if role != constant.InterconnectRoleMaster {
		c.JSON(403, gin.H{"error": "interconnect master role not enabled on this server"})
		return
	}

	var node models.InterconnectNode
	if err := database.DB.Where("token = ?", tokenStr).First(&node).Error; err != nil {
		c.JSON(401, gin.H{"error": "invalid interconnect token"})
		return
	}

	nodeID := node.ID

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.Errorf("[Tunnel] 升级隧道连接失败 (NodeID: %s): %v", nodeID, err)
		return
	}

	// 将 WebSocket 包装为标准的 net.Conn
	netConn := NetConn(conn)

	// 获取客户端真实 IP 和端口以更新连接地址，方便面板显示
	clientIP := c.ClientIP()
	port := "unknown"
	if parts := strings.Split(netConn.RemoteAddr().String(), ":"); len(parts) > 1 {
		port = parts[len(parts)-1]
	}
	realURL := "tunnel://" + clientIP + ":" + port

	if node.URL != realURL {
		database.DB.Model(&node).Update("url", realURL)
	}

	// 我们是主节点（父节点），主动发起对子节点的请求。
	// 因此在 Yamux 协议中我们扮演 Client，而子节点扮演 Server。
	conf := yamux.DefaultConfig()
	conf.EnableKeepAlive = true

	session, err := yamux.Client(netConn, conf)
	if err != nil {
		logger.Errorf("[Tunnel] 启动 Yamux Client 失败 (NodeID: %s): %v", nodeID, err)
		netConn.Close()
		return
	}

	AddSession(nodeID, tokenStr, session)
	logger.Infof("[Tunnel] 已建立来自子节点的隧道连接: %s (%s)", nodeID, node.Name)

	go func(s *yamux.Session) {
		<-s.CloseChan()
		logger.Infof("[Tunnel] 子节点隧道连接已断开 (NodeID: %s, IsClosed: %v)", nodeID, s.IsClosed())
		RemoveSession(nodeID, s)
	}(session)
}