package a2a

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"agent-squad/squad"
)

// conformance_test.go — Q3 WP4 验收：双向 conformance（本地假 A2A 服务端）。
// 发布侧：BuildCard + CardHandler 在 well-known 路径服务合法 Agent Card；
// 消费侧：FetchCard 发现校验 + Executor 走完任务状态机
// （submitted→working→completed / failed / input-required），终态对齐
// squad.ExecResult。

// fakeA2AServer 最小 A2A 1.0 服务端：well-known 卡片 + JSON-RPC
// （message/send 建任务；tasks/get 按脚本推进状态）。
type fakeA2AServer struct {
	mu     sync.Mutex
	card   AgentCard
	tasks  map[string]*Task
	plan   []string // 每次 tasks/get 推进到的状态序列（走完停在末态）
	calls  map[string]int
	artifacts []Artifact
	statusMsg string
}

func newFakeServer(t *testing.T, card AgentCard, plan []string, artifacts []Artifact) *fakeA2AServer {
	t.Helper()
	s := &fakeA2AServer{card: card, tasks: map[string]*Task{}, plan: plan,
		calls: map[string]int{}, artifacts: artifacts}
	return s
}

func (s *fakeA2AServer) handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(WellKnownPath, func(w http.ResponseWriter, r *http.Request) {
		s.mu.Lock()
		body, _ := json.Marshal(s.card)
		s.mu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write(body)
	})
	mux.HandleFunc("/rpc", func(w http.ResponseWriter, r *http.Request) {
		var req rpcRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeRPCError(w, 0, -32700, "parse error")
			return
		}
		s.mu.Lock()
		defer s.mu.Unlock()
		s.calls[req.Method]++
		switch req.Method {
		case "message/send":
			id := fmt.Sprintf("task_%d", len(s.tasks)+1)
			state := StateSubmitted
			if len(s.plan) > 0 {
				state = s.plan[0]
			}
			task := &Task{Kind: "task", ID: id, ContextID: "ctx_1",
				Status: TaskStatus{State: state, Message: s.statusMessage()}}
			s.tasks[id] = task
			writeRPCResult(w, req.ID, task)
		case "tasks/get":
			var params struct {
				ID string `json:"id"`
			}
			raw, _ := json.Marshal(req.Params)
			_ = json.Unmarshal(raw, &params)
			task, ok := s.tasks[params.ID]
			if !ok {
				writeRPCError(w, req.ID, -32001, "task not found")
				return
			}
			// 按脚本推进状态机（停留在终态）。
			if !isTerminal(task.Status.State) && len(s.plan) > 0 {
				n := s.calls["tasks/get#"+params.ID]
				s.calls["tasks/get#"+params.ID] = n + 1
				idx := n + 1
				if idx >= len(s.plan) {
					idx = len(s.plan) - 1
				}
				task.Status.State = s.plan[idx]
				task.Status.Message = s.statusMessage()
				if task.Status.State == StateCompleted {
					task.Artifacts = s.artifacts
				}
			}
			writeRPCResult(w, req.ID, task)
		default:
			writeRPCError(w, req.ID, -32601, "method not found")
		}
	})
	return mux
}

func (s *fakeA2AServer) statusMessage() *Message {
	if s.statusMsg == "" {
		return nil
	}
	return &Message{Kind: "message", Role: "agent",
		Parts: []Part{{Kind: "text", Text: s.statusMsg}}}
}

func writeRPCResult(w http.ResponseWriter, id int64, result any) {
	body, _ := json.Marshal(result)
	env, _ := json.Marshal(rpcResponse{JSONRPC: "2.0", ID: id, Result: body})
	w.Header().Set("Content-Type", "application/json")
	_, _ = w.Write(env)
}

func writeRPCError(w http.ResponseWriter, id int64, code int, msg string) {
	env, _ := json.Marshal(rpcResponse{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: code, Message: msg}})
	w.Header().Set("Content-Type", "application/json")
	_, _ = w.Write(env)
}

// ---- 发布侧 ----

func TestBuildCardFromRegistries(t *testing.T) {
	reg := squad.NewRoleRegistry()
	for _, r := range []squad.RoleSpec{
		{Role: "researcher", Name: "调研员", Keywords: []string{"调研", "检索"},
			Instructions: "围绕目标调研关键问题。\n第二行不进摘要。"},
		{Role: "writer", Name: "撰稿员", Keywords: []string{"撰写"}},
	} {
		if err := reg.Register(r); err != nil {
			t.Fatal(err)
		}
	}
	bps := squad.NewBlueprintRegistry()
	if err := bps.Register(squad.Blueprint{Key: "brief", Title: "简报",
		Description: "生成简报", Match: []string{"简报", "brief"},
		Steps: []squad.StepSpec{{ID: "s1", Role: "writer", Title: "成文"}}}); err != nil {
		t.Fatal(err)
	}
	card, err := BuildCard(reg, bps, CardOptions{
		Name: "示例小队", Description: "演示团队能力", URL: "https://team.example/rpc",
	})
	if err != nil {
		t.Fatal(err)
	}
	if card.ProtocolVersion != ProtocolVersion || card.Version != "1.0.0" {
		t.Fatalf("card = %+v", card)
	}
	byID := map[string]AgentSkill{}
	for _, s := range card.Skills {
		byID[s.ID] = s
	}
	if byID["researcher"].Name != "调研员" || byID["researcher"].Description != "围绕目标调研关键问题。" {
		t.Fatalf("role skill = %+v", byID["researcher"])
	}
	if len(byID["researcher"].Tags) != 2 {
		t.Fatalf("tags = %v", byID["researcher"].Tags)
	}
	if byID["brief"].Name != "简报" || strings.Join(byID["brief"].Examples, ",") != "简报,brief" {
		t.Fatalf("blueprint skill = %+v", byID["brief"])
	}
	if _, err := BuildCard(reg, bps, CardOptions{}); err == nil {
		t.Fatal("缺 URL 必须 fail-closed")
	}
}

// CardHandler 以 JSON 服务卡片（发布端点形态）。
func TestCardHandlerServesJSON(t *testing.T) {
	card := AgentCard{ProtocolVersion: ProtocolVersion, Name: "小队", URL: "https://x/rpc", Version: "1.0.0"}
	ts := httptest.NewServer(CardHandler(card))
	defer ts.Close()
	resp, err := http.Get(ts.URL + WellKnownPath)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if ct := resp.Header.Get("Content-Type"); ct != "application/json" {
		t.Fatalf("content-type = %s", ct)
	}
	var got AgentCard
	if err := json.NewDecoder(resp.Body).Decode(&got); err != nil {
		t.Fatal(err)
	}
	if got.Name != "小队" || got.ProtocolVersion != ProtocolVersion {
		t.Fatalf("card = %+v", got)
	}
}

// ---- 消费侧：发现 + 状态机 ----

func TestFetchCardAndExecuteCompleted(t *testing.T) {
	fileBytes := base64.StdEncoding.EncodeToString([]byte("notes 内容"))
	artifacts := []Artifact{{
		ArtifactID: "art1", Name: "调研结果",
		Parts: []Part{
			{Kind: "text", Text: "# 调研结论\n\n结论正文，长度足够作为交付物内容。"},
			{Kind: "file", File: &FilePart{Name: "notes.md", Bytes: fileBytes, MimeType: "text/markdown"}},
		},
	}}
	card := AgentCard{ProtocolVersion: ProtocolVersion, Name: "远端小队",
		URL: "PLACEHOLDER", Version: "1.0.0"}
	srv := newFakeServer(t, card, []string{StateWorking, StateWorking, StateCompleted}, artifacts)
	ts := httptest.NewServer(srv.handler())
	defer ts.Close()
	srv.mu.Lock()
	srv.card.URL = ts.URL + "/rpc"
	srv.mu.Unlock()

	// 发现：well-known 路径取卡并校验。
	got, err := FetchCard(ts.Client(), ts.URL+WellKnownPath)
	if err != nil {
		t.Fatal(err)
	}
	if got.Name != "远端小队" || got.URL != ts.URL+"/rpc" {
		t.Fatalf("card = %+v", got)
	}

	// 执行：submitted/working → completed，工件映射 patches。
	var states []string
	ex := NewExecutor(got)
	ex.Client = ts.Client()
	ex.PollInterval = 5 * time.Millisecond
	ex.OnState = func(s string) { states = append(states, s) }
	res, err := ex.Execute(context.Background(), squad.ExecRequest{
		Prompt: "调研主题 X", Targets: []string{"output.md"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if !res.Success {
		t.Fatalf("result = %+v", res)
	}
	if string(res.Patches["notes.md"]) != "notes 内容" {
		t.Fatalf("file part 未落成 patch：%q", res.Patches["notes.md"])
	}
	if !strings.Contains(string(res.Patches["output.md"]), "调研结论") {
		t.Fatalf("文本部件未写入目标：%q", res.Patches["output.md"])
	}
	joined := strings.Join(states, ",")
	if !strings.Contains(joined, StateWorking) || !strings.HasSuffix(joined, StateCompleted) {
		t.Fatalf("状态机迁移 = %v", states)
	}
	srv.mu.Lock()
	sends := srv.calls["message/send"]
	srv.mu.Unlock()
	if sends != 1 {
		t.Fatalf("message/send = %d", sends)
	}
}

func TestExecuteMapsFailedAndInputRequired(t *testing.T) {
	card := AgentCard{ProtocolVersion: ProtocolVersion, Name: "远端", URL: "PLACEHOLDER", Version: "1"}
	srv := newFakeServer(t, card, []string{StateWorking, StateFailed}, nil)
	srv.statusMsg = "远端执行崩溃"
	ts := httptest.NewServer(srv.handler())
	defer ts.Close()
	srv.mu.Lock()
	srv.card.URL = ts.URL + "/rpc"
	srv.mu.Unlock()

	ex := NewExecutor(card)
	ex.BaseURL = ts.URL + "/rpc"
	ex.Client = ts.Client()
	ex.PollInterval = 5 * time.Millisecond
	res, err := ex.Execute(context.Background(), squad.ExecRequest{Prompt: "x", Targets: []string{"output.md"}})
	if err != nil {
		t.Fatal(err)
	}
	if res.Success || !strings.Contains(res.ErrorReason, "failed") || !strings.Contains(res.ErrorReason, "远端执行崩溃") {
		t.Fatalf("failed 映射 = %+v", res)
	}

	// input-required：fail-closed（跨进程人工输入不在适配面内）。
	srv2 := newFakeServer(t, card, []string{StateInputRequired}, nil)
	srv2.statusMsg = "请补充授权"
	ts2 := httptest.NewServer(srv2.handler())
	defer ts2.Close()
	ex2 := &Executor{BaseURL: ts2.URL + "/rpc", Client: ts2.Client(), PollInterval: 5 * time.Millisecond}
	res2, err := ex2.Execute(context.Background(), squad.ExecRequest{Prompt: "x", Targets: []string{"output.md"}})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Success || !strings.Contains(res2.ErrorReason, "input-required") {
		t.Fatalf("input-required 映射 = %+v", res2)
	}
}

// 直接回 Message（无任务态）→ 文本落目标。
func TestExecuteDirectMessageResult(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req rpcRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		writeRPCResult(w, req.ID, Message{Kind: "message", Role: "agent",
			Parts: []Part{{Kind: "text", Text: "直接回答正文"}}})
	}))
	defer ts.Close()
	ex := &Executor{BaseURL: ts.URL, Client: ts.Client()}
	res, err := ex.Execute(context.Background(), squad.ExecRequest{Prompt: "x", Targets: []string{"output.md"}})
	if err != nil {
		t.Fatal(err)
	}
	if !res.Success || !strings.Contains(string(res.Patches["output.md"]), "直接回答正文") {
		t.Fatalf("res = %+v", res)
	}
}

// 传输/协议错误 → error（引擎按瞬态评估），不吞成成功。
func TestExecuteProtocolErrors(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeRPCError(w, 1, -32601, "method not found")
	}))
	defer ts.Close()
	ex := &Executor{BaseURL: ts.URL, Client: ts.Client()}
	if _, err := ex.Execute(context.Background(), squad.ExecRequest{Prompt: "x"}); err == nil ||
		!strings.Contains(err.Error(), "协议错误") {
		t.Fatalf("err = %v", err)
	}
	// ctx 取消传播。
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	slow := newFakeServer(t, AgentCard{ProtocolVersion: "1.0", Name: "x", URL: ts.URL, Version: "1"},
		[]string{StateWorking}, nil)
	ts2 := httptest.NewServer(slow.handler())
	defer ts2.Close()
	ex2 := &Executor{BaseURL: ts2.URL + "/rpc", Client: ts2.Client(), PollInterval: time.Second}
	if _, err := ex2.Execute(ctx, squad.ExecRequest{Prompt: "x"}); err == nil {
		t.Fatal("取消必须传播")
	}
}

// FetchCard fail-closed：缺必备字段拒绝。
func TestFetchCardValidates(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"protocolVersion":"1.0","name":"","url":"","version":""}`))
	}))
	defer ts.Close()
	if _, err := FetchCard(ts.Client(), ts.URL); err == nil || !strings.Contains(err.Error(), "name") {
		t.Fatalf("err = %v", err)
	}
}