package handler

import (
	"context"
	"encoding/json"
	"fkteams/agentcore"
	"fkteams/engine"
	"fkteams/events"
	"fkteams/events/log"
	"fkteams/g"
	"fkteams/server/handler/taskstream"
	"fkteams/server/origin"
	"fkteams/tools/approval"
	"fkteams/tools/ask"
	"log"
	"sync"
	"time"

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

var upgrader = websocket.Upgrader{
	CheckOrigin:     origin.IsAllowed,
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// WSMessage WebSocket 消息类型
type WSMessage struct {
	Type        string        `json:"type"`
	SessionID   string        `json:"session_id,omitempty"`
	Offset      uint64        `json:"offset,omitempty"`
	Message     string        `json:"message,omitempty"`
	Mode        string        `json:"mode,omitempty"`
	AgentName   string        `json:"agent_name,omitempty"`
	Decision    int           `json:"decision,omitempty"`
	Contents    []ContentPart `json:"contents,omitempty"`
	AskSelected []string      `json:"ask_selected,omitempty"`
	AskFreeText string        `json:"ask_free_text,omitempty"`
}

// WebSocketHandler 处理 WebSocket 连接
func WebSocketHandler() gin.HandlerFunc {
	return func(c *gin.Context) {
		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			log.Printf("websocket upgrade failed: %v", err)
			return
		}

		connCtx, connCancel := context.WithCancel(c.Request.Context())
		registerConn(conn, connCancel)
		sm := getSessionManager(conn)

		defer func() {
			connCancel()
			removeSessionManager(conn)
			unregisterConn(conn)
			_ = conn.Close()
		}()

		// 监听 context 取消，主动关闭连接
		go func() {
			<-connCtx.Done()
			_ = conn.Close()
		}()

		// 提取 WS 连接时的 token，用于每条消息的二次校验
		wsToken := c.Query("token")
		if wsToken == "" {
			if cookie, err := c.Cookie("fk_token"); err == nil {
				wsToken = cookie
			}
		}
		authEnabled, _ := AuthEnabled()

		// 线程安全的写入
		var writeMu sync.Mutex
		writeJSON := func(v any) error {
			writeMu.Lock()
			defer writeMu.Unlock()
			return conn.WriteJSON(v)
		}

		_ = writeJSON(map[string]any{
			"type":    events.NotifyConnected,
			"message": "欢迎连接到非空小队",
		})

		for {
			_, msgBytes, err := conn.ReadMessage()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					log.Printf("websocket read error: %v", err)
				}
				break
			}

			var wsMsg WSMessage
			if err := json.Unmarshal(msgBytes, &wsMsg); err != nil {
				_ = writeJSON(map[string]any{"type": events.NotifyError, "error": "invalid message format"})
				continue
			}

			// chat/steer 消息二次校验 token，防止 WS 长时间连接后 token 过期仍可操作
			if (wsMsg.Type == "chat" || wsMsg.Type == "follow_up" || wsMsg.Type == "steer" || wsMsg.Type == "steering") && authEnabled && !ValidateToken(wsToken) {
				_ = writeJSON(map[string]any{"type": events.NotifyError, "error": "登录已过期，请重新登录"})
				conn.Close()
				break
			}

			switch wsMsg.Type {
			case "chat", "follow_up":
				if wsMsg.SessionID == "" {
					_ = writeJSON(map[string]any{"type": events.NotifyError, "error": "session_id is required"})
					continue
				}
				go handleChatMessage(sm, wsMsg, writeJSON)

			case "steer", "steering":
				handleSteeringMessage(wsMsg, writeJSON)

			case "resume":
				sid := wsMsg.SessionID
				if sid == "" {
					_ = writeJSON(map[string]any{"type": events.NotifyError, "error": "session_id is required"})
					continue
				}
				stream := GlobalStreams.Get(sid)
				if stream != nil {
					ok, subID := stream.Subscribe(taskstream.FuncSubscriber(writeJSON), wsMsg.Offset)
					if ok {
						// 成功重新绑定并回放事件
						sm.mu.Lock()
						sm.tasks[sid] = &sessionTask{cancel: stream.Cancel, stream: stream, subID: subID}
						sm.mu.Unlock()
						log.Printf("task resumed: session=%s", sid)
					} else {
						_ = writeJSON(map[string]any{
							"type":       events.NotifyProcessingEnd,
							"session_id": sid,
							"message":    "任务已完成或不存在",
						})
					}
				} else {
					_ = writeJSON(map[string]any{
						"type":       events.NotifyProcessingEnd,
						"session_id": sid,
						"message":    "任务已完成或不存在",
					})
				}

			case "cancel":
				sid := wsMsg.SessionID
				if sid != "" {
					GlobalStreams.CancelAndRemove(sid)
					sm.cancelTask(sid)
				}
				_ = writeJSON(map[string]any{"type": "cancellation_requested", "session_id": sid, "message": "取消请求已发送"})

			case "approval":
				sid := wsMsg.SessionID
				if stream := GlobalStreams.Get(sid); stream != nil {
					_ = stream.SubmitInterrupt(taskstream.InterruptApproval, wsMsg.Decision)
				}

			case "ask_response":
				sid := wsMsg.SessionID
				resp := &ask.AskResponse{
					Selected: wsMsg.AskSelected,
					FreeText: wsMsg.AskFreeText,
				}
				if stream := GlobalStreams.Get(sid); stream != nil {
					_ = stream.SubmitInterrupt(taskstream.InterruptAsk, resp)
				}

			case "ping":
				_ = writeJSON(map[string]any{"type": events.NotifyPong})

			default:
				_ = writeJSON(map[string]any{"type": events.NotifyError, "error": "unknown message type"})
			}
		}
	}
}

// --- WebSocket HITL 中断处理器 ---

// buildInterruptHandler 构建 WebSocket 聊天的 HITL 中断处理器
func buildInterruptHandler(recorder *eventlog.HistoryRecorder, sessionID string, writeJSON func(any) error, stream *taskstream.Stream) engine.InterruptHandler {
	channelHandler := engine.ChannelHandler(stream.InterruptCh())
	return func(ctx context.Context, interrupts []agentcore.Interrupt) (map[string]any, error) {
		// 检查是否为 ask_questions 中断
		if info := extractAskInfo(interrupts); info != nil {
			stream.BeginInterrupt(taskstream.InterruptAsk)
			defer stream.CompleteInterrupt(taskstream.InterruptAsk)

			recorder.RecordEvent(events.Event{
				Type:       events.EventAction,
				ActionType: events.ActionAskQuestions,
				Content:    info.Question,
			})
			payload := map[string]any{
				"type":         events.NotifyAskQuestions,
				"session_id":   sessionID,
				"question":     info.Question,
				"options":      info.Options,
				"multi_select": info.MultiSelect,
			}
			_ = writeJSON(payload)

			result, err := channelHandler(ctx, interrupts)

			if err == nil {
				recorder.RecordEvent(events.Event{
					Type:       events.EventAction,
					ActionType: events.ActionAskResponse,
					Content:    askResponseText(result),
				})
			}
			return result, err
		}

		// 默认审批流程
		msg := extractInterruptMessage(interrupts)

		stream.BeginInterrupt(taskstream.InterruptApproval)
		defer stream.CompleteInterrupt(taskstream.InterruptApproval)

		recorder.RecordEvent(events.Event{
			Type:       events.EventAction,
			ActionType: events.ActionApprovalRequired,
			Content:    msg,
		})
		_ = writeJSON(map[string]any{
			"type":       events.NotifyApprovalRequired,
			"session_id": sessionID,
			"message":    msg,
		})

		result, err := channelHandler(ctx, interrupts)

		if err == nil {
			if text := approvalDecisionText(result); text != "" {
				recorder.RecordEvent(events.Event{
					Type:       events.EventAction,
					ActionType: events.ActionApprovalDecision,
					Content:    text,
				})
			}
		}

		return result, err
	}
}

// --- WebSocket 事件回调 ---

// wsEventCallbackBuffered 构建支持断线缓冲的事件回调
func wsEventCallbackBuffered(recorder *eventlog.HistoryRecorder, sessionID string, stream *taskstream.Stream) func(events.Event) error {
	return func(event events.Event) error {
		if event.Type == events.EventAction && event.ActionType == events.ActionInterrupted {
			return nil
		}
		recorder.RecordEvent(event)
		data := convertEventToMap(event)
		data["session_id"] = sessionID
		stream.Publish(data)
		return nil
	}
}

// --- WebSocket 聊天处理 ---

// handleChatMessage 处理 WebSocket 聊天消息
func handleChatMessage(sm *sessionManager, wsMsg WSMessage, writeJSON func(any) error) {
	sessionID := wsMsg.SessionID
	mode := wsMsg.Mode
	if mode == "" {
		mode = "team"
	}
	if wsMsg.Message == "" && len(wsMsg.Contents) == 0 {
		_ = writeJSON(map[string]any{
			"type":       events.NotifyError,
			"session_id": sessionID,
			"error":      "message or contents is required",
		})
		return
	}

	if existing := GlobalStreams.Get(sessionID); existing != nil && existing.Status() == "processing" {
		enqueueTaskMessage(existing, sessionID, taskstream.QueueFollowUp, wsMsg.Message, wsMsg.Contents)
		return
	}

	// 任务 context 独立于连接——断连不会自动取消任务
	taskCtx, taskCancel := context.WithCancel(context.Background())
	defer taskCancel()

	// 注册到统一 TaskStream（支持断线重连 + Push/Pull 消费）
	stream := GlobalStreams.Register(taskstream.StreamConfig{
		SessionID:  sessionID,
		Cancel:     taskCancel,
		CleanupTTL: 5 * time.Minute,
	})
	// 绑定当前 WS 连接为 Push 订阅者
	_, subID := stream.Subscribe(taskstream.FuncSubscriber(writeJSON), 0)
	defer func() {
		stream.Done()
	}()

	// 同时在连接级 sessionManager 中注册（用于取消和断线 Unsubscribe）
	taskID := sm.startTask(sessionID, taskCancel)
	sm.mu.Lock()
	if t, exists := sm.tasks[sessionID]; exists && t.id == taskID {
		t.stream = stream
		t.subID = subID
	}
	sm.mu.Unlock()
	defer sm.removeTask(sessionID, taskID)

	// 获取 runner
	r, err := resolveRunner(taskCtx, mode, wsMsg.AgentName)
	if err != nil {
		log.Printf("failed to resolve runner: session=%s, err=%v", sessionID, err)
		stream.Publish(map[string]any{"type": events.NotifyError, "session_id": sessionID, "error": err.Error()})
		return
	}

	// 构建输入消息
	recorder := eventlog.GlobalSessionManager.GetOrCreate(sessionID, historyDir)
	turnInput, userDisplayText := buildChatInput(recorder, wsMsg.Message, wsMsg.Contents)
	stream.Publish(map[string]any{
		"type":       events.NotifyUserMessage,
		"session_id": sessionID,
		"content":    userDisplayText,
	})

	publishFn := func(v any) error { stream.Publish(v.(map[string]any)); return nil }
	interruptHandler := buildInterruptHandler(recorder, sessionID, publishFn, stream)
	steeringSource := buildSteeringSource(stream, recorder, sessionID)
	currentInput := turnInput
	currentDisplayText := userDisplayText
	updateSessionTitleAndStatus(sessionID, currentDisplayText, "processing")
	stream.Publish(map[string]any{
		"type":       events.NotifyProcessingStart,
		"session_id": sessionID,
		"message":    "开始处理您的请求...",
	})

	for {
		_, runErr := engine.NewSession(r, sessionID).
			WithInput(currentInput).
			OnEvent(wsEventCallbackBuffered(recorder, sessionID, stream)).
			WithHistory(recorder).
			OnInterrupt(interruptHandler).
			NonInteractive().
			WithContext(approval.RegistryContext(approval.NewDefaultRegistry())).
			WithContext(func(ctx context.Context) context.Context {
				return agentcore.WithSteeringSource(ctx, steeringSource)
			}).
			Run(taskCtx)
		if runErr != nil {
			if isConnectionClosed(taskCtx, runErr) {
				log.Printf("task cancelled: session=%s", sessionID)
				stream.SetStatus("cancelled")
				stream.Publish(map[string]any{
					"type":       events.NotifyCancelled,
					"session_id": sessionID,
					"message":    "任务已取消",
				})
				finishCancelledChat(recorder, sessionID, currentDisplayText)
				return
			}
			log.Printf("failed to run task: session=%s, err=%v", sessionID, runErr)
			stream.SetStatus("error")
			stream.Publish(map[string]any{"type": events.NotifyError, "session_id": sessionID, "error": runErr.Error()})
			finishErrorChat(recorder, sessionID, currentDisplayText, runErr)
			return
		}

		recorder.FinalizeCurrent()
		saveHistory(recorder, chatHistoryPath(sessionID), sessionID)
		if queued, ok := stream.DequeueNextMessage(); ok {
			publishQueueUpdated(stream, sessionID)
			currentDisplayText = queued.DisplayText
			currentInput = buildQueuedChatInput(recorder, queued)
			updateSessionTitleAndStatus(sessionID, currentDisplayText, "processing")
			publishQueuedExecutionStart(stream, sessionID, queued)
			continue
		}

		stream.SetStatus("completed")
		stream.Publish(map[string]any{
			"type":       events.NotifyProcessingEnd,
			"session_id": sessionID,
			"message":    "处理完成",
		})
		ensureSessionMetadataWithStatus(sessionID, currentDisplayText, "completed")
		if g.MemoryManager != nil {
			g.MemoryManager.ExtractFromRecorder(recorder, sessionID)
		}
		return
	}
}

func handleSteeringMessage(wsMsg WSMessage, writeJSON func(any) error) {
	sessionID := wsMsg.SessionID
	if sessionID == "" {
		_ = writeJSON(map[string]any{"type": events.NotifyError, "error": "session_id is required"})
		return
	}
	if wsMsg.Message == "" && len(wsMsg.Contents) == 0 {
		_ = writeJSON(map[string]any{
			"type":       events.NotifyError,
			"session_id": sessionID,
			"error":      "message or contents is required",
		})
		return
	}
	stream := GlobalStreams.Get(sessionID)
	if stream == nil || stream.Status() != "processing" {
		_ = writeJSON(map[string]any{
			"type":       events.NotifyError,
			"session_id": sessionID,
			"error":      "no running task to steer",
		})
		return
	}
	enqueueTaskMessage(stream, sessionID, taskstream.QueueSteering, wsMsg.Message, wsMsg.Contents)
}
