package server

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

	"agent/internal/agent"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 4096,
	// 本地单用户 Web，允许任意 Origin
	CheckOrigin: func(r *http.Request) bool { return true },
}

// wsInbound 客户端 → 服务端消息。
// chat：在绑定会话上运行 Agent；switch：切换绑定会话；stop：停止当前生成；ping：保活。
type wsInbound struct {
	Type      string `json:"type"`
	Content   string `json:"content"`
	SessionID string `json:"session_id"`
}

// wsConn 封装一个 WebSocket 连接：绑定一个会话、串行处理消息。
type wsConn struct {
	server *Server
	conn   *websocket.Conn
	sendMu sync.Mutex // 串行化 WriteJSON

	sessionID string
	unreg     func() // 当前会话的 live 注销函数

	runMu     sync.Mutex
	runCancel context.CancelFunc
}

func (s *Server) ws(c *gin.Context) {
	sessionID := strings.TrimSpace(c.Query("session"))
	if sessionID == "" || !sessionIDPattern.MatchString(sessionID) {
		log.Printf("ws: 连接被拒（session 参数无效）session=%q", sessionID)
		fail(c, http.StatusBadRequest, "缺少有效的 session 参数")
		return
	}
	if _, ok := s.store.Get(sessionID); !ok {
		log.Printf("ws: 连接被拒（会话不存在）session=%q", sessionID)
		fail(c, http.StatusNotFound, "会话不存在")
		return
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	client := &wsConn{server: s, conn: conn}
	client.bind(sessionID)
	client.readPump()
}

// sendJSON 安全地推送一个事件帧 {type, data}。
func (cl *wsConn) sendJSON(eventType string, data any) {
	cl.sendMu.Lock()
	defer cl.sendMu.Unlock()
	_ = cl.conn.WriteJSON(map[string]any{"type": eventType, "data": data})
}

// bind 切换绑定会话：注销旧会话的 live 注册，注册新会话。
// 传入空串表示仅注销。
func (cl *wsConn) bind(sessionID string) {
	if cl.unreg != nil {
		cl.unreg()
		cl.unreg = nil
	}
	cl.sessionID = sessionID
	if sessionID == "" {
		return
	}
	cl.unreg = cl.server.agent.RegisterLive(sessionID, func(event agent.Event) {
		cl.sendJSON(event.Type, event.Data)
	})
}

func (cl *wsConn) handleChat(content string) {
	content = strings.TrimSpace(content)
	if content == "" {
		cl.sendJSON("error", map[string]any{"message": "消息内容不能为空"})
		return
	}
	// 未配置 LLM API Key 时直接提示去配置，避免进入无效的模型调用。
	// 环境变量 LLM_API_KEY（托管模式）与设置页 api-config.json 都计入已配置。
	if !cl.server.apiCfg.IsAPIKeyConfigured() {
		cl.sendJSON("error", map[string]any{"message": "尚未配置 LLM API，请先点击左下角设置按钮配置 API Key 与模型"})
		return
	}
	cl.runMu.Lock()
	if cl.runCancel != nil {
		cl.runMu.Unlock()
		cl.sendJSON("error", map[string]any{"message": "Agent 正在执行，请等待完成或先停止"})
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	cl.runCancel = cancel
	cl.runMu.Unlock()

	done := make(chan struct{})
	go func() {
		defer close(done)
		_ = cl.server.agent.Run(ctx, cl.sessionID, content, func(event agent.Event) {
			cl.sendJSON(event.Type, event.Data)
		})
	}()
	go func() {
		<-done
		cl.runMu.Lock()
		cl.runCancel = nil
		cl.runMu.Unlock()
	}()
}

func (cl *wsConn) handleStop() {
	cl.runMu.Lock()
	if cl.runCancel != nil {
		cl.runCancel()
	}
	cl.runMu.Unlock()
}

func (cl *wsConn) readPump() {
	defer func() {
		cl.bind("") // 注销 live
		_ = cl.conn.Close()
	}()
	cl.conn.SetReadLimit(1 << 20) // 1MB
	_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
	cl.conn.SetPongHandler(func(string) error {
		_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
		return nil
	})

	for {
		// 任何收到消息（含客户端心跳 {type:"ping"}）都刷新 90s 读超时，
		// 否则浏览器只能发协议层 Pong，应用层 ping 不会触发 SetPongHandler，连接必死在第 90 秒
		_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
		_, raw, err := cl.conn.ReadMessage()
		if err != nil {
			return
		}
		var msg wsInbound
		if err := json.Unmarshal(raw, &msg); err != nil {
			continue
		}
		switch msg.Type {
		case "chat":
			cl.handleChat(msg.Content)
		case "stop":
			cl.handleStop()
		case "switch":
			if msg.SessionID != "" && sessionIDPattern.MatchString(msg.SessionID) {
				if _, ok := cl.server.store.Get(msg.SessionID); ok {
					cl.bind(msg.SessionID)
				}
			}
		case "ping":
			cl.sendJSON("pong", map[string]any{"time": time.Now().UnixMilli()})
		default:
			log.Printf("ws: unknown message type %q", msg.Type)
		}
	}
}