package node

import (
	"net"
	"sync"
	"sync/atomic"

	"git.zeroonesoft.cn/golib/zonat/agent"

	"github.com/xtaci/smux"
)

// Session 一个 agent 连接的会话：smux 服务端 + 控制流处理。
// 数据流由节点侧隧道监听主动 OpenStream；agent 主动打开的第一条流是控制流，
// 用于注册/注销隧道与心跳。
type Session struct {
	agentId string
	node    *Node
	session *smux.Session

	control   net.Conn   // agent 打开的控制流
	controlMu sync.Mutex // 控制流写锁（handleControl 与 pushClosed 并发写）
	dead      atomic.Bool
}

func (s *Session) isDead() bool { return s.dead.Load() }

func (s *Session) close() {
	s.dead.Store(true)
	s.session.Close()
}

// OpenStream 建立一条到 agent 的数据流
func (s *Session) OpenStream() (net.Conn, error) {
	return s.session.OpenStream()
}

// loop 接受 agent 打开的控制流，直到会话结束
func (s *Session) loop() {
	for {
		stream, err := s.session.AcceptStream()
		if err != nil {
			s.close()
			return
		}
		// agent 打开的第一条流是控制流，pushClosed 依赖它；
		// 赋值与写入共用 controlMu，避免与推送协程竞争
		s.controlMu.Lock()
		if s.control == nil {
			s.control = stream
		}
		s.controlMu.Unlock()
		go s.handleControl(stream)
	}
}

// handleControl 处理控制流上的帧：注册/注销隧道、心跳。
// 单条控制流内的帧顺序处理；应答与节点推送共用 controlMu 串行写入。
func (s *Session) handleControl(stream net.Conn) {
	defer stream.Close()
	for {
		f, err := agent.ReadFrame(stream)
		if err != nil {
			return
		}
		switch f.Cmd() {
		case agent.CmdRegisterTunnel:
			t := agent.Tunnel{}
			ret := agent.Ret{Code: 1, Msg: "注册失败"}
			if err := f.Unmarshal(&t); err != nil {
				ret.Msg = err.Error()
			} else if nt, err := s.node.registerTunnel(s, t); err != nil {
				ret.Msg = err.Error()
			} else {
				ret = agent.Ret{Code: 0, Msg: "注册成功", Port: nt.ListenPort}
			}
			s.reply(stream, f, ret)
		case agent.CmdUnregisterTunnel:
			t := agent.Tunnel{}
			ret := agent.Ret{Code: 1, Msg: "注销失败"}
			if err := f.Unmarshal(&t); err != nil {
				ret.Msg = err.Error()
			} else if s.node.unregisterTunnel(t.Id) {
				ret = agent.Ret{Code: 0, Msg: "注销成功"}
			} else {
				ret.Msg = "隧道不存在"
			}
			s.reply(stream, f, ret)
		case agent.CmdPing:
			s.writeControl(func() error {
				return agent.WriteFrame(stream, agent.NewFrame(agent.FrameVersion, agent.CmdPing, f.StreamID()))
			})
		case agent.CmdLogin, agent.CmdTarget, agent.CmdTunnelClosed:
			// 不该出现在 agent→节点方向，忽略
		default:
			s.reply(stream, f, agent.Ret{Code: 1, Msg: "未知命令"})
		}
	}
}

func (s *Session) reply(stream net.Conn, req agent.Frame, ret agent.Ret) {
	s.writeControl(func() error {
		f := agent.NewFrame(agent.FrameVersion, req.Cmd(), req.StreamID())
		if err := f.Marshal(ret); err != nil {
			return err
		}
		return agent.WriteFrame(stream, f)
	})
}

// pushClosed 节点主动通知 agent：隧道已被删除（TTL 过期等）
func (s *Session) pushClosed(tunnelId string) {
	if s.control == nil {
		return
	}
	s.writeControl(func() error {
		f := agent.NewFrame(agent.FrameVersion, agent.CmdTunnelClosed, 0)
		if err := f.Marshal(agent.Tunnel{Id: tunnelId}); err != nil {
			return err
		}
		return agent.WriteFrame(s.control, f)
	})
}

func (s *Session) writeControl(fn func() error) {
	s.controlMu.Lock()
	defer s.controlMu.Unlock()
	if s.dead.Load() {
		return
	}
	// 写失败意味着会话已死，由读循环感知并收尾
	_ = fn()
}

func smuxServer(conn net.Conn) (*smux.Session, error) {
	return smux.Server(conn, agent.SmuxConfig())
}