// Agent Card
public record AgentCard(
    String name,
    String description,
    String url,
    String version,
    Capabilities capabilities,
    Authentication authentication,
    List<String> defaultInputModes,
    List<String> defaultOutputModes,
    List<Skill> skills
) {
    public record Capabilities(boolean streaming, boolean pushNotifications, boolean stateTransitionHistory) {}
    public record Authentication(List<String> schemes) {}
    public record Skill(String id, String name, String description, List<String> tags, List<String> examples) {}
}

// Task
public record Task(
    String id,
    TaskStatus status,
    Message message,
    List<Artifact> artifacts,
    Instant createdAt,
    Instant updatedAt
) {
    public enum TaskStatus {
        SUBMITTED, WORKING, COMPLETED, FAILED, CANCELED
    }
}

// Message
public record Message(String role, List<Part> parts) {}

// Part - sealed interface for type safety
public sealed interface Part {
    record TextPart(String type, String text) implements Part {}
    record FilePart(String type, FileInfo file) implements Part {
        public record FileInfo(String name, String mimeType, String bytes) {}
    }
    record DataPart(String type, Map<String, Object> data) implements Part {}
}

// Artifact
public record Artifact(String name, List<Part> parts) {}

@RestController
public class A2aController {

    private final DocumentAnalysisService analysisService;
    private final Map<String, Task> taskStore = new ConcurrentHashMap<>();

    @GetMapping("/.well-known/agent.json")
    public AgentCard getAgentCard() {
        return new AgentCard(
            "Document Analysis Agent",
            "文档分析 Agent，支持文本摘要、关键词提取、情感分析",
            "http://localhost:8080",
            "1.0.0",
            new AgentCard.Capabilities(true, false, true),
            new AgentCard.Authentication(List.of("Bearer")),
            List.of("text", "file"),
            List.of("text", "data"),
            List.of(
                new AgentCard.Skill(
                    "summarize", "文档摘要",
                    "对长文档进行智能摘要，提取核心信息",
                    List.of("nlp", "summary"),
                    List.of("帮我总结一下这份报告的核心内容")
                ),
                new AgentCard.Skill(
                    "extract-keywords", "关键词提取",
                    "从文档中提取关键词和主题",
                    List.of("nlp", "keywords"),
                    List.of("提取这篇文章的关键词")
                )
            )
        );
    }

    // JSON-RPC 端点
    @PostMapping("/a2a")
    public Object handleJsonRpc(@RequestBody JsonRpcRequest request) {
        return switch (request.method()) {
            case "tasks/send" -> handleSend(request);
            case "tasks/get" -> handleGet(request);
            case "tasks/cancel" -> handleCancel(request);
            default -> new JsonRpcError(request.id(), -32601, "Method not found");
        };
    }
}

private JsonRpcResponse handleSend(JsonRpcRequest request) {
    var params = (Map<String, Object>) request.params();
    var taskId = (String) params.get("id");
    var messageMap = (Map<String, Object>) params.get("message");
    var parts = parseParts((List<Map<String, Object>>) messageMap.get("parts"));

    // 创建 Task（状态：SUBMITTED）
    var task = new Task(taskId, Task.TaskStatus.SUBMITTED,
        new Message("user", parts), new ArrayList<>(), Instant.now(), Instant.now());
    taskStore.put(taskId, task);

    // 异步处理
    CompletableFuture.runAsync(() -> {
        try {
            // 更新状态为 WORKING
            updateTaskStatus(taskId, Task.TaskStatus.WORKING);

            // 执行分析
            var result = analysisService.analyze(parts);

            // 更新为 COMPLETED，附加 Artifact
            var taskResult = taskStore.get(taskId);
            taskResult.artifacts().add(result);
            updateTaskStatus(taskId, Task.TaskStatus.COMPLETED);
        } catch (Exception e) {
            updateTaskStatus(taskId, Task.TaskStatus.FAILED);
        }
    });

    return new JsonRpcResponse(request.id(), task);
}

@Service
public class DocumentAnalysisService {

    private final ChatModel chatModel;

    public DocumentAnalysisService(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    public Artifact analyze(List<Part> parts) {
        // 提取文本内容
        String textContent = parts.stream()
            .filter(p -> p instanceof Part.TextPart)
            .map(p -> ((Part.TextPart) p).text())
            .collect(Collectors.joining("\n"));

        // 使用 Spring AI 进行分析
        var systemPrompt = """
            你是一个专业的文档分析助手。请对以下文档进行分析，返回：
            1. 核心摘要（100字以内）
            2. 关键词列表（5-10个）
            3. 情感倾向（正面/中性/负面）
            4. 文档类型判断

            以 JSON 格式返回结果。
            """;

        var response = chatModel.call(
            new Prompt(List.of(
                new SystemMessage(systemPrompt),
                new UserMessage(textContent)
            ))
        );

        var analysisResult = parseAnalysisResult(response.getResult().getOutput().getText());

        return new Artifact("文档分析结果", List.of(
            new Part.DataPart("data", analysisResult)
        ));
    }
}

@GetMapping(value = "/a2a/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public Flux<ServerSentEvent<Object>> streamTask(@RequestParam String taskId) {
    return Flux.create(sink -> {
        // 注册状态变更监听
        var listener = new TaskEventListener() {
            @Override
            public void onStatusChange(Task task) {
                sink.next(ServerSentEvent.builder()
                    .event("task-status")
                    .data(task)
                    .build());

                if (task.status() == Task.TaskStatus.COMPLETED
                    || task.status() == Task.TaskStatus.FAILED) {
                    sink.complete();
                }
            }

            @Override
            public void onArtifact(Task task, Artifact artifact) {
                sink.next(ServerSentEvent.builder()
                    .event("task-artifact")
                    .data(artifact)
                    .build());
            }
        };

        taskEventBus.register(taskId, listener);
        sink.onDispose(() -> taskEventBus.unregister(taskId, listener));
    });
}

@Service
public class A2aClientAgent {

    private final RestTemplate restTemplate;
    private final ChatModel chatModel;

    // 1. 发现目标 Agent 的能力
    public AgentCard discoverAgent(String agentUrl) {
        return restTemplate.getForObject(
            agentUrl + "/.well-known/agent.json", AgentCard.class);
    }

    // 2. 发送任务
    public Task sendTask(String agentUrl, String taskId, String message) {
        var request = Map.of(
            "jsonrpc", "2.0",
            "id", UUID.randomUUID().toString(),
            "method", "tasks/send",
            "params", Map.of(
                "id", taskId,
                "message", Map.of(
                    "role", "user",
                    "parts", List.of(Map.of("type", "text", "text", message))
                )
            )
        );

        var response = restTemplate.postForObject(
            agentUrl + "/a2a", request, JsonRpcResponse.class);

        return (Task) response.result();
    }

    // 3. 轮询任务状态
    public Task pollTask(String agentUrl, String taskId) {
        var request = Map.of(
            "jsonrpc", "2.0",
            "id", UUID.randomUUID().toString(),
            "method", "tasks/get",
            "params", Map.of("id", taskId)
        );

        var response = restTemplate.postForObject(
            agentUrl + "/a2a", request, JsonRpcResponse.class);

        return (Task) response.result();
    }

    // 4. 完整的委托流程
    public String delegateToAgent(String agentUrl, String userRequest) {
        // 发现能力
        var card = discoverAgent(agentUrl);
        log.info("发现 Agent: {} - {}", card.name(), card.description());

        // 创建任务
        String taskId = UUID.randomUUID().toString();
        var task = sendTask(agentUrl, taskId, userRequest);

        // 等待完成（带超时）
        Instant deadline = Instant.now().plus(Duration.ofMinutes(5));
        while (Instant.now().isBefore(deadline)) {
            task = pollTask(agentUrl, taskId);

            if (task.status() == Task.TaskStatus.COMPLETED) {
                return extractResult(task);
            } else if (task.status() == Task.TaskStatus.FAILED) {
                throw new RuntimeException("Agent task failed: " + taskId);
            }

            Thread.sleep(Duration.ofSeconds(2));
        }

        throw new RuntimeException("Agent task timeout: " + taskId);
    }
}

@RestController
@RequestMapping("/customer-service")
public class CustomerServiceController {

    private final A2aClientAgent a2aClient;
    private final ChatModel chatModel;

    @Value("${agents.document-analysis.url}")
    private String documentAnalysisAgentUrl;

    @PostMapping("/handle")
    public String handleCustomerRequest(@RequestBody CustomerRequest request) {
        // Step 1: 用 LLM 理解用户意图
        var intent = understandIntent(request.message());

        // Step 2: 根据意图决定是否需要委托给其他 Agent
        if (intent.needsDocumentAnalysis()) {
            // 通过 A2A 委托给文档分析 Agent
            var analysisResult = a2aClient.delegateToAgent(
                documentAnalysisAgentUrl,
                request.message()
            );

            // Step 3: 基于分析结果生成回复
            return generateResponse(request, analysisResult);
        }

        return generateDirectResponse(request);
    }
}

@Configuration
public class A2aSecurityConfig {

    @Bean
    public SecurityFilterChain a2aSecurityChain(HttpSecurity http) throws Exception {
        return http
            .securityMatcher("/a2a/**", "/.well-known/agent.json")
            .oauth2ResourceServer(oauth2 -> oauth2
                .jwt(Customizer.withDefaults()))
            .authorizeHttpRequests(auth -> auth
                .requestMatchers("/.well-known/agent.json").permitAll()
                .requestMatchers("/a2a/**").authenticated()
                .anyRequest().denyAll())
            .build();
    }
}

@Configuration
public class A2aClientConfig {

    @Bean
    public RestTemplate a2aRestTemplate() {
        var factory = new SimpleClientHttpRequestFactory();
        factory.setConnectTimeout(Duration.ofSeconds(5));
        factory.setReadTimeout(Duration.ofMinutes(5));  // Agent 任务可能很慢
        return new RestTemplate(factory);
    }
}

@Bean
public ObservationRegistryCustomizer<ObservationRegistry> a2aObservation() {
    return registry -> registry.observationConfig()
        .observationHandler(new A2aObservationHandler());
}

public JsonRpcResponse handleSend(JsonRpcRequest request) {
    String taskId = extractTaskId(request);

    // 幂等检查：如果任务已存在且非终态，返回现有任务
    Task existing = taskStore.get(taskId);
    if (existing != null && !isTerminal(existing.status())) {
        return new JsonRpcResponse(request.id(), existing);
    }

    // ... 创建新任务
}