package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/history"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/task"
)

// 本文件锁定 design/11 §7.2 —— `viewer`（只读监控角色）过去经由 /ws/chat
// 越界的最后两处，都是实现 §7.1 下行门时发现的：
//
//	A. 握手时的历史回放（ws_history.go）携带主会话的消息正文，因此仅限
//	   owner/admin；
//	B. 每一个上行（客户端→服务端）JSON 帧（ws_route.go）对非 owner/admin
//	   连接默认拒，因此 viewer 可以跟随会话状态机，但永远发不出指令。
//
// 与 §7.1 下行门一起，它们钉住了合并后的契约：
// `viewer` = 可以连接、可以读它那 12 个只读 REST 监控端点、可以接收
// 7 个状态机信号 —— 除此之外在 WS 上收不到也发不出任何东西。这里的连接与
// 令牌都是真的（httptest server + 签名 JWT），所以路由表里的错误无法藏在
// 假实现后面。
const (
	// wsRoleGateReadWait 界定一次预期会产出帧的读取。gorilla 在读超时后会
	// 弄坏连接，因此这里每条连接至多做一次超时读取，且总是作为最后一次读。
	wsRoleGateReadWait = 3 * time.Second
	// wsRoleGateNoFrameWait 是「本不该存在的帧」若存在就会到达的窗口
	// （定向回复由处理函数同步入队，所以这么长的空转窗口即已足够判定）。
	wsRoleGateNoFrameWait = 250 * time.Millisecond
	// wsRoleGateSettle 让被接受帧的副作用先落地，再断言它们不存在。
	wsRoleGateSettle = 200 * time.Millisecond
	// wsRoleGateWriteWait 界定一次控制帧写入（ping / close）。
	wsRoleGateWriteWait = 2 * time.Second
)

// countingRunner 是带调用计数的 Chatter 接缝：design/11 §7.2 B 必须证明
// viewer 的 `chat` 帧不会到达模型，所以「runner 是否被驱动过」要可观测，
// 而不能从广播反推。
type countingRunner struct {
	mu    sync.Mutex
	calls int
}

func (r *countingRunner) Run(context.Context, string, []model.Message) (*loop.Result, error) {
	r.mu.Lock()
	r.calls++
	r.mu.Unlock()
	return &loop.Result{Text: "reply", Steps: 1, StopReason: "stop"}, nil
}

func (r *countingRunner) count() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.calls
}

// replayHistoryStore 用固定转录内容提供握手回放，同时继承 fakeHistory 的
// Append 记录，因此一个 store 既让回放确定，又能报告被接受的 `chat` 帧持久化了
// 什么。
type replayHistoryStore struct {
	*fakeHistory
	mu   sync.Mutex
	msgs []history.Message
}

func newReplayHistoryStore(msgs ...history.Message) *replayHistoryStore {
	return &replayHistoryStore{fakeHistory: &fakeHistory{}, msgs: msgs}
}

func (s *replayHistoryStore) Recent(context.Context, string, int) ([]history.Message, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]history.Message(nil), s.msgs...), nil
}

// historyFrameWire 是解码后用于断言的握手回放帧。
type historyFrameWire struct {
	Type      string `json:"type"`
	SessionID string `json:"sessionId"`
	Data      struct {
		Messages []struct {
			Role     string `json:"role"`
			Content  string `json:"content"`
			Thinking string `json:"thinking"`
		} `json:"messages"`
	} `json:"data"`
}

// errorFrameWire 是解码后用于断言的任意 {type:"error"} 帧。
type errorFrameWire struct {
	Type string `json:"type"`
	Data struct {
		Code    string `json:"code"`
		Message string `json:"message"`
	} `json:"data"`
}

// roleGateEnv 为这些测试接线一个 gateway：带计数的 runner（使泄漏的上行帧
// 表现为一次模型调用）、会回放转录并记录持久化的 store，以及子代理/任务的
// 接缝，使泄漏的 subagent_ask 或 task_output_query 会产生一个真实且可断言的
// 内容帧，而不是处理函数级错误。
type roleGateEnv struct {
	srv  *Server
	url  string
	jwt  *auth.JWTService
	run  *countingRunner
	hist *replayHistoryStore
}

func newRoleGateEnv(t *testing.T, transcript []history.Message) *roleGateEnv {
	t.Helper()
	run := &countingRunner{}
	hist := newReplayHistoryStore(transcript...)
	srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{
		Loop:    run,
		History: hist,
		Scheduler: &stubAuditor{res: &orchestration.AuditResult{
			AgentID: "sub-1",
			Mode:    orchestration.AuditFull,
			Goal:    "子代理的私密目标",
		}},
		Tasks: &stubTaskOutputs{ev: task.OutputEvent{
			ID:     "task-1",
			Output: task.Output{Stdout: "任务的私密输出"},
		}},
	})
	return &roleGateEnv{srv: srv, url: ts.URL, jwt: jwt, run: run, hist: hist}
}

func (e *roleGateEnv) token(t *testing.T, role auth.Role) string {
	t.Helper()
	return signAny(t, e.jwt, "acct", role)
}

// TestWSHandshakeReplayPerRole 锁定 design/11 §7.2 A：握手历史回放持有主会话的
// 消息正文，因此只有 owner/admin 连接会收到它。泄漏的回放永远是连接上的第一个
// 帧 —— 它在 upgrade 期间就已入队，早于读泵启动、也早于任何回复 —— 所以读
// viewer 的第一帧就把「没有回放」变成精确断言，而不是靠猜一个空转窗口。
func TestWSHandshakeReplayPerRole(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, []history.Message{
		{ID: "m1", SessionID: history.MainSessionID, Role: history.RoleUser, Content: "主人的私密提问", Thinking: "不该泄漏的推理", CreatedAt: 1},
		{ID: "m2", SessionID: history.MainSessionID, Role: history.RoleAssistant, Content: "角色的私密回答", CreatedAt: 2},
	})

	owner := dialWS(t, env.url, env.token(t, auth.RoleOwner))
	defer owner.Close()
	admin := dialWS(t, env.url, env.token(t, auth.RoleAdmin))
	defer admin.Close()
	viewer := dialWS(t, env.url, env.token(t, auth.RoleViewer))
	defer viewer.Close()
	waitConnCount(t, env.srv, 3)

	for _, tc := range []struct {
		role auth.Role
		conn *websocket.Conn
	}{{auth.RoleOwner, owner}, {auth.RoleAdmin, admin}} {
		typ, raw := readWire(t, tc.conn, string(tc.role)+" handshake")
		if typ != "history" {
			t.Fatalf("%s first frame = %q, want the history replay", tc.role, typ)
		}
		var hf historyFrameWire
		if err := json.Unmarshal(raw, &hf); err != nil {
			t.Fatalf("%s decode history frame %s: %v", tc.role, raw, err)
		}
		if hf.SessionID != history.MainSessionID {
			t.Fatalf("%s history sessionId = %q, want %q", tc.role, hf.SessionID, history.MainSessionID)
		}
		if len(hf.Data.Messages) != 2 {
			t.Fatalf("%s replayed %d messages, want 2 (%+v)", tc.role, len(hf.Data.Messages), hf.Data.Messages)
		}
		if got := hf.Data.Messages[0].Content; got != "主人的私密提问" {
			t.Fatalf("%s replayed content = %q, want the stored user message", tc.role, got)
		}
		for i, m := range hf.Data.Messages {
			if m.Thinking != "" {
				t.Fatalf("%s replayed message[%d] carries thinking %q; it must be stripped", tc.role, i, m.Thinking)
			}
		}
	}

	// 这里主动探一次 viewer，以保证一定有帧回来（§7.2 B 的拒绝）；若回放泄漏，
	// 它会排在那条回复之前。
	if err := viewer.WriteJSON(wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"viewer"}`)}); err != nil {
		t.Fatalf("viewer write chat: %v", err)
	}
	if typ, raw := readWire(t, viewer, "viewer handshake"); typ == "history" {
		t.Fatalf("viewer received the handshake replay %s; design/11 §7.2 A allows owner/admin only", raw)
	}
}

// TestWSInboundDeniedForViewer 锁定 design/11 §7.2 B：非 owner/admin 连接不得
// 发送任何应用内指令。每个被拒的帧都得到定向到该连接的
// {type:"error", data:{code:"forbidden"}}，别的什么都没有 —— 没有内容帧、
// 没有广播、没有模型调用、没有持久化消息。
func TestWSInboundDeniedForViewer(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, nil)

	// 第二条 viewer 连接用于观察拒绝是指向发送方的：若拒绝走广播，这里也会收到。
	observer := dialWS(t, env.url, env.token(t, auth.RoleViewer))
	defer observer.Close()

	cases := []struct {
		name  string
		frame wsInbound
	}{
		{"chat", wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"替主人发一条"}`)}},
		{"interrupt", wsInbound{Type: "interrupt", Data: json.RawMessage(`{"sessionId":"web"}`)}},
		{"subagent_ask", wsInbound{Type: "subagent_ask", Data: json.RawMessage(`{"id":"sub-1","mode":"full"}`)}},
		{"task_output_query", wsInbound{Type: "task_output_query", Data: json.RawMessage(`{"taskId":"task-1"}`)}},
		{"voice_input", wsInbound{Type: "voice_input", Data: json.RawMessage(`{"sessionId":"web","text":"你好"}`)}},
		{"mode_switch", wsInbound{Type: "mode_switch", Data: json.RawMessage(`{"mode":"chat"}`)}},
	}
	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			conn := dialWS(t, env.url, env.token(t, auth.RoleViewer))
			defer conn.Close()
			if err := conn.WriteJSON(tc.frame); err != nil {
				t.Fatalf("write %s: %v", tc.name, err)
			}
			typ, raw := readWire(t, conn, tc.name)
			if typ != "error" {
				t.Fatalf("%s reply = %q (%s), want a directed error frame", tc.name, typ, raw)
			}
			var ef errorFrameWire
			if err := json.Unmarshal(raw, &ef); err != nil {
				t.Fatalf("%s decode reply %s: %v", tc.name, raw, err)
			}
			if ef.Data.Code != wsForbiddenCode {
				t.Fatalf("%s error code = %q, want %q (frame=%s)", tc.name, ef.Data.Code, wsForbiddenCode, raw)
			}
			// 对 subagent_ask / task_output_query 而言，泄漏会是一条定向的
			// subagent_event / task_output 回复，携带子代理上下文或任务输出，
			// 因此断言拒绝就是全部应答。
			assertNoFrame(t, conn, tc.name+" follow-up")
		})
	}

	// `chat` 这一例正是过去会驱动模型并持久化消息的那个；给它（已被拒的）
	// 副作用一点时间，确认它们没有发生。
	time.Sleep(wsRoleGateSettle)
	calls, records := env.run.count(), env.hist.records()
	if calls != 0 || len(records) != 0 {
		t.Fatalf("viewer frames reached the runner %d time(s) and persisted %d message(s) %+v; a viewer may do neither", calls, len(records), records)
	}
	assertNoFrame(t, observer, "observer viewer (a refusal must not be broadcast)")
}

// TestWSInboundChatDrivesOwnerTurn 是 design/11 §7.2 B 的主人侧对照：被拒给 viewer
// 的那个帧（`chat`）对主人依然会驱动模型并持久化回合，从而证明这道门确实是角色
// 门，而不是 chat 路径坏了。
func TestWSInboundChatDrivesOwnerTurn(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, nil)
	owner := dialWS(t, env.url, env.token(t, auth.RoleOwner))
	defer owner.Close()
	waitConnCount(t, env.srv, 1)

	if err := owner.WriteJSON(wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"主人问一句"}`)}); err != nil {
		t.Fatalf("write chat: %v", err)
	}
	waitUntil(t, "the owner's chat to drive the runner", func() bool { return env.run.count() == 1 })
	waitUntil(t, "the owner's turn to be persisted", func() bool { return len(env.hist.records()) == 2 })

	got := env.hist.records()
	if got[0].Role != history.RoleUser || got[1].Role != history.RoleAssistant {
		t.Fatalf("persisted roles = %q, %q, want user then assistant", got[0].Role, got[1].Role)
	}
	if got[1].Content != storedParagraph("reply") {
		t.Fatalf("persisted reply = %q, want the runner's text in its canonical (paragraph-wrapped) form", got[1].Content)
	}
}

// TestWSInboundGateKeepsProtocolFrames 证明上行门不会杀死连接
// （design/11 §7.2 B「协议级帧不受影响」）：被拒一次之后，viewer 仍能完成
// ping/pong 往返和正常关闭，服务端也会注销它。
func TestWSInboundGateKeepsProtocolFrames(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, nil)
	viewer := dialWS(t, env.url, env.token(t, auth.RoleViewer))
	waitConnCount(t, env.srv, 1)

	if err := viewer.WriteJSON(wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"x"}`)}); err != nil {
		t.Fatalf("write chat: %v", err)
	}
	if typ, raw := readWire(t, viewer, "probe"); typ != "error" {
		t.Fatalf("probe reply = %q (%s), want the forbidden error frame", typ, raw)
	}

	// 第二次被拒的往返证明读泵在第一轮之后依然存活：门在任何处理函数之前就返回，
	// 连接未被触动。
	if err := viewer.WriteJSON(wsInbound{Type: "interrupt", Data: json.RawMessage(`{"sessionId":"web"}`)}); err != nil {
		t.Fatalf("write interrupt after refusal: %v", err)
	}
	if typ, raw := readWire(t, viewer, "second probe"); typ != "error" {
		t.Fatalf("second probe reply = %q (%s), want the forbidden error frame", typ, raw)
	}

	// 客户端 ping 由服务端默认的 ping 处理函数应答。pong 会被下面的客户端读取
	// 消费掉，因此由该处理函数来报告它。这是本连接唯一的一次超时读取
	// （gorilla 在读截止时间过后会弄坏连接），所以它排在读侧所有断言之后。
	ponged := make(chan struct{}, 1)
	viewer.SetPongHandler(func(string) error {
		select {
		case ponged <- struct{}{}:
		default:
		}
		return nil
	})
	if err := viewer.WriteControl(websocket.PingMessage, nil, time.Now().Add(wsRoleGateWriteWait)); err != nil {
		t.Fatalf("write ping: %v", err)
	}
	_ = viewer.SetReadDeadline(time.Now().Add(wsRoleGateReadWait))
	if _, msg, err := viewer.ReadMessage(); err == nil {
		t.Fatalf("data frame %s while waiting for the pong", msg)
	}
	select {
	case <-ponged:
	default:
		t.Fatal("the server never answered the ping; the gate broke the protocol path")
	}

	// 正常关闭依然会注销该连接。
	if err := viewer.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(wsRoleGateWriteWait)); err != nil {
		t.Fatalf("write close: %v", err)
	}
	waitConnCount(t, env.srv, 0)
}

// readWire 读取一个帧，返回它的类型与原始字节（以便调用方解码按类型区分的
// 载荷）。若在截止时间前没有帧到达则让测试失败。
func readWire(t *testing.T, conn *websocket.Conn, what string) (string, []byte) {
	t.Helper()
	_ = conn.SetReadDeadline(time.Now().Add(wsRoleGateReadWait))
	_, msg, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("%s: read frame: %v", what, err)
	}
	var env struct {
		Type string `json:"type"`
	}
	if err := json.Unmarshal(msg, &env); err != nil {
		t.Fatalf("%s: decode frame %s: %v", what, msg, err)
	}
	return env.Type, msg
}

// assertNoFrame 断言短窗口内 conn 上不会再有任何东西到达。它必须是该连接的最后
// 一次读取：读截止时间到期后 gorilla 会留下一个已损坏的连接。
func assertNoFrame(t *testing.T, conn *websocket.Conn, what string) {
	t.Helper()
	_ = conn.SetReadDeadline(time.Now().Add(wsRoleGateNoFrameWait))
	_, msg, err := conn.ReadMessage()
	if err == nil {
		t.Fatalf("%s: unexpected frame %s", what, msg)
	}
	var netErr net.Error
	if !errors.As(err, &netErr) || !netErr.Timeout() {
		t.Fatalf("%s: read failed with %v, want a timeout (no frame expected)", what, err)
	}
}

// waitUntil 轮询 cond 直到它成立或超过截止时间，使测试能对异步回合（run → persist）
// 做断言，而不必用一段可能空跑通过的固定 sleep。
func waitUntil(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(wsRoleGateReadWait)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(2 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}