/*
 * Copyright 2024-2024 the original author or authors.
 */

package io.modelcontextprotocol.client;

import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.fail;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Function;

import io.modelcontextprotocol.spec.McpSchema.ElicitFormRequest;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpError;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.BlobResourceContents;
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
import io.modelcontextprotocol.spec.McpSchema.ElicitRequest;
import io.modelcontextprotocol.spec.McpSchema.ElicitResult;
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
import io.modelcontextprotocol.spec.McpSchema.Prompt;
import io.modelcontextprotocol.spec.McpSchema.ReadResourceResult;
import io.modelcontextprotocol.spec.McpSchema.Resource;
import io.modelcontextprotocol.spec.McpSchema.ResourceContents;
import io.modelcontextprotocol.spec.McpSchema.Root;
import io.modelcontextprotocol.spec.McpSchema.SubscribeRequest;
import io.modelcontextprotocol.spec.McpSchema.TextResourceContents;
import io.modelcontextprotocol.spec.McpSchema.Tool;
import io.modelcontextprotocol.spec.McpSchema.UnsubscribeRequest;
import io.modelcontextprotocol.spec.McpTransport;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.test.StepVerifier;

/**
 * Test suite for the {@link McpAsyncClient} that can be used with different
 * {@link McpTransport} implementations.
 *
 * @author Christian Tzolov
 * @author Dariusz Jędrzejczyk
 */
public abstract class AbstractMcpAsyncClientTests {

	private static final String ECHO_TEST_MESSAGE = "Hello MCP Spring AI!";

	abstract protected McpClientTransport createMcpTransport();

	protected Duration getRequestTimeout() {
		return Duration.ofSeconds(14);
	}

	protected Duration getInitializationTimeout() {
		return Duration.ofSeconds(20);
	}

	McpAsyncClient client(McpClientTransport transport) {
		return client(transport, Function.identity());
	}

	McpAsyncClient client(McpClientTransport transport, Function<McpClient.AsyncSpec, McpClient.AsyncSpec> customizer) {
		AtomicReference<McpAsyncClient> client = new AtomicReference<>();

		assertThatCode(() -> {
			McpClient.AsyncSpec builder = McpClient.async(transport)
				.requestTimeout(getRequestTimeout())
				.initializationTimeout(getInitializationTimeout())
				.sampling(req -> Mono.just(CreateMessageResult
					.builder(McpSchema.Role.USER, McpSchema.TextContent.builder("Oh, hi!").build(), "modelId")
					.stopReason(CreateMessageResult.StopReason.END_TURN)
					.build()))
				.capabilities(ClientCapabilities.builder().roots(true).sampling().build());
			builder = customizer.apply(builder);
			client.set(builder.build());
		}).doesNotThrowAnyException();

		return client.get();
	}

	void withClient(McpClientTransport transport, Consumer<McpAsyncClient> c) {
		withClient(transport, Function.identity(), c);
	}

	void withClient(McpClientTransport transport, Function<McpClient.AsyncSpec, McpClient.AsyncSpec> customizer,
			Consumer<McpAsyncClient> c) {
		var client = client(transport, customizer);
		try {
			c.accept(client);
		}
		finally {
			StepVerifier.create(client.closeGracefully()).expectComplete().verify(Duration.ofSeconds(10));
		}
	}

	<T> void verifyNotificationSucceedsWithImplicitInitialization(Function<McpAsyncClient, Mono<T>> operation,
			String action) {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(operation.apply(mcpAsyncClient)).verifyComplete();
		});
	}

	<T> void verifyCallSucceedsWithImplicitInitialization(Function<McpAsyncClient, Mono<T>> operation, String action) {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(operation.apply(mcpAsyncClient)).expectNextCount(1).verifyComplete();
		});
	}

	@Test
	void testConstructorWithInvalidArguments() {
		assertThatThrownBy(() -> McpClient.async(null).build()).isInstanceOf(IllegalArgumentException.class)
			.hasMessage("Transport must not be null");

		assertThatThrownBy(() -> McpClient.async(createMcpTransport()).requestTimeout(null).build())
			.isInstanceOf(IllegalArgumentException.class)
			.hasMessage("Request timeout must not be null");
	}

	@Test
	void testListToolsWithoutInitialization() {
		verifyCallSucceedsWithImplicitInitialization(client -> client.listTools(McpSchema.FIRST_PAGE), "listing tools");
	}

	@Test
	void testListTools() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listTools(McpSchema.FIRST_PAGE)))
				.consumeNextWith(result -> {
					assertThat(result.tools()).isNotNull().isNotEmpty();

					Tool firstTool = result.tools().get(0);
					assertThat(firstTool.name()).isNotNull();
					assertThat(firstTool.description()).isNotNull();
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllTools() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listTools()))
				.consumeNextWith(result -> {
					assertThat(result.tools()).isNotNull().isNotEmpty();

					Tool firstTool = result.tools().get(0);
					assertThat(firstTool.name()).isNotNull();
					assertThat(firstTool.description()).isNotNull();
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllToolsReturnsImmutableList() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listTools()))
				.consumeNextWith(result -> {
					assertThat(result.tools()).isNotNull();
					// Verify that the returned list is immutable
					assertThatThrownBy(() -> result.tools()
						.add(Tool.builder("test", JSON_MAPPER, "{\"type\":\"object\"}").title("test").build()))
						.isInstanceOf(UnsupportedOperationException.class);
				})
				.verifyComplete();
		});
	}

	@Test
	void testPingWithoutInitialization() {
		verifyCallSucceedsWithImplicitInitialization(client -> client.ping(), "pinging the server");
	}

	@Test
	void testPing() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.ping()))
				.expectNextCount(1)
				.verifyComplete();
		});
	}

	@Test
	void testCallToolWithoutInitialization() {
		CallToolRequest callToolRequest = CallToolRequest.builder("echo")
			.arguments(Map.of("message", ECHO_TEST_MESSAGE))
			.build();
		verifyCallSucceedsWithImplicitInitialization(client -> client.callTool(callToolRequest), "calling tools");
	}

	@Test
	void testCallTool() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			CallToolRequest callToolRequest = CallToolRequest.builder("echo")
				.arguments(Map.of("message", ECHO_TEST_MESSAGE))
				.build();

			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.callTool(callToolRequest)))
				.consumeNextWith(callToolResult -> {
					assertThat(callToolResult).isNotNull().satisfies(result -> {
						assertThat(result.content()).isNotNull();
						assertThat(result.isError()).isNull();
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testCallToolWithInvalidTool() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			CallToolRequest invalidRequest = CallToolRequest.builder("nonexistent_tool")
				.arguments(Map.of("message", ECHO_TEST_MESSAGE))
				.build();

			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.callTool(invalidRequest)))
				.consumeErrorWith(
						e -> assertThat(e).isInstanceOf(McpError.class).hasMessage("Unknown tool: nonexistent_tool"))
				.verify();
		});
	}

	@ParameterizedTest
	@ValueSource(strings = { "success", "error", "debug" })
	void testCallToolWithMessageAnnotations(String messageType) {
		McpClientTransport transport = createMcpTransport();

		withClient(transport, mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize()
				.then(mcpAsyncClient.callTool(McpSchema.CallToolRequest.builder("annotatedMessage")
					.arguments(Map.of("messageType", messageType, "includeImage", true))
					.build())))
				.consumeNextWith(result -> {
					assertThat(result).isNotNull();
					assertThat(result.isError()).isNotEqualTo(true);
					assertThat(result.content()).isNotEmpty();
					assertThat(result.content()).allSatisfy(content -> {
						switch (content.type()) {
							case "text":
								McpSchema.TextContent textContent = assertInstanceOf(McpSchema.TextContent.class,
										content);
								assertThat(textContent.text()).isNotEmpty();
								assertThat(textContent.annotations()).isNotNull();

								switch (messageType) {
									case "error":
										assertThat(textContent.annotations().priority()).isEqualTo(1.0);
										assertThat(textContent.annotations().audience())
											.containsOnly(McpSchema.Role.USER, McpSchema.Role.ASSISTANT);
										break;
									case "success":
										assertThat(textContent.annotations().priority()).isEqualTo(0.7);
										assertThat(textContent.annotations().audience())
											.containsExactly(McpSchema.Role.USER);
										break;
									case "debug":
										assertThat(textContent.annotations().priority()).isEqualTo(0.3);
										assertThat(textContent.annotations().audience())
											.containsExactly(McpSchema.Role.ASSISTANT);
										break;
									default:
										throw new IllegalStateException("Unexpected value: " + content.type());
								}
								break;
							case "image":
								McpSchema.ImageContent imageContent = assertInstanceOf(McpSchema.ImageContent.class,
										content);
								assertThat(imageContent.data()).isNotEmpty();
								assertThat(imageContent.annotations()).isNotNull();
								assertThat(imageContent.annotations().priority()).isEqualTo(0.5);
								assertThat(imageContent.annotations().audience()).containsExactly(McpSchema.Role.USER);
								break;
							default:
								fail("Unexpected content type: " + content.type());
						}
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testListResourcesWithoutInitialization() {
		verifyCallSucceedsWithImplicitInitialization(client -> client.listResources(McpSchema.FIRST_PAGE),
				"listing resources");
	}

	@Test
	void testListResources() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listResources(McpSchema.FIRST_PAGE)))
				.consumeNextWith(resources -> {
					assertThat(resources).isNotNull().satisfies(result -> {
						assertThat(result.resources()).isNotNull();

						if (!result.resources().isEmpty()) {
							Resource firstResource = result.resources().get(0);
							assertThat(firstResource.uri()).isNotNull();
							assertThat(firstResource.name()).isNotNull();
						}
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllResources() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listResources()))
				.consumeNextWith(resources -> {
					assertThat(resources).isNotNull().satisfies(result -> {
						assertThat(result.resources()).isNotNull();

						if (!result.resources().isEmpty()) {
							Resource firstResource = result.resources().get(0);
							assertThat(firstResource.uri()).isNotNull();
							assertThat(firstResource.name()).isNotNull();
						}
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllResourcesReturnsImmutableList() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listResources()))
				.consumeNextWith(result -> {
					assertThat(result.resources()).isNotNull();
					// Verify that the returned list is immutable
					assertThatThrownBy(() -> result.resources().add(Resource.builder("test://uri", "test").build()))
						.isInstanceOf(UnsupportedOperationException.class);
				})
				.verifyComplete();
		});
	}

	@Test
	void testMcpAsyncClientState() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			assertThat(mcpAsyncClient).isNotNull();
		});
	}

	@Test
	void testListPromptsWithoutInitialization() {
		verifyCallSucceedsWithImplicitInitialization(client -> client.listPrompts(McpSchema.FIRST_PAGE),
				"listing " + "prompts");
	}

	@Test
	void testListPrompts() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listPrompts(McpSchema.FIRST_PAGE)))
				.consumeNextWith(prompts -> {
					assertThat(prompts).isNotNull().satisfies(result -> {
						assertThat(result.prompts()).isNotNull();

						if (!result.prompts().isEmpty()) {
							Prompt firstPrompt = result.prompts().get(0);
							assertThat(firstPrompt.name()).isNotNull();
							assertThat(firstPrompt.description()).isNotNull();
						}
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllPrompts() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listPrompts()))
				.consumeNextWith(prompts -> {
					assertThat(prompts).isNotNull().satisfies(result -> {
						assertThat(result.prompts()).isNotNull();

						if (!result.prompts().isEmpty()) {
							Prompt firstPrompt = result.prompts().get(0);
							assertThat(firstPrompt.name()).isNotNull();
							assertThat(firstPrompt.description()).isNotNull();
						}
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllPromptsReturnsImmutableList() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listPrompts()))
				.consumeNextWith(result -> {
					assertThat(result.prompts()).isNotNull();
					// Verify that the returned list is immutable
					assertThatThrownBy(() -> result.prompts()
						.add(Prompt.builder("test").title("test").description("test").build()))
						.isInstanceOf(UnsupportedOperationException.class);
				})
				.verifyComplete();
		});
	}

	@Test
	void testGetPromptWithoutInitialization() {
		GetPromptRequest request = GetPromptRequest.builder("simple_prompt").arguments(Map.of()).build();
		verifyCallSucceedsWithImplicitInitialization(client -> client.getPrompt(request), "getting " + "prompts");
	}

	@Test
	void testGetPrompt() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier
				.create(mcpAsyncClient.initialize()
					.then(mcpAsyncClient
						.getPrompt(GetPromptRequest.builder("simple_prompt").arguments(Map.of()).build())))
				.consumeNextWith(prompt -> {
					assertThat(prompt).isNotNull().satisfies(result -> {
						assertThat(result.messages()).isNotEmpty();
						assertThat(result.messages()).hasSize(1);
					});
				})
				.verifyComplete();
		});
	}

	@Test
	void testRootsListChangedWithoutInitialization() {
		verifyNotificationSucceedsWithImplicitInitialization(client -> client.rootsListChangedNotification(),
				"sending roots list changed notification");
	}

	@Test
	void testRootsListChanged() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.rootsListChangedNotification()))
				.verifyComplete();
		});
	}

	@Test
	void testInitializeWithRootsListProviders() {
		withClient(createMcpTransport(),
				builder -> builder.roots(Root.builder("file:///test/path").name("test-root").build()), client -> {
					StepVerifier.create(client.initialize()).expectNextCount(1).verifyComplete();
				});
	}

	@Test
	void testAddRoot() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			Root newRoot = Root.builder("file:///new/test/path").name("new-test-root").build();
			StepVerifier.create(mcpAsyncClient.addRoot(newRoot)).verifyComplete();
		});
	}

	@Test
	void testAddRootWithNullValue() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.addRoot(null))
				.consumeErrorWith(e -> assertThat(e).isInstanceOf(IllegalArgumentException.class)
					.hasMessage("Root must not be null"))
				.verify();
		});
	}

	@Test
	void testRemoveRoot() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			Root root = Root.builder("file:///test/path/to/remove").name("root-to-remove").build();
			StepVerifier.create(mcpAsyncClient.addRoot(root)).verifyComplete();

			StepVerifier.create(mcpAsyncClient.removeRoot(root.uri())).verifyComplete();
		});
	}

	@Test
	void testRemoveNonExistentRoot() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.removeRoot("nonexistent-uri"))
				.consumeErrorWith(e -> assertThat(e).isInstanceOf(IllegalStateException.class)
					.hasMessage("Root with uri 'nonexistent-uri' not found"))
				.verify();
		});
	}

	@Test
	void testReadResource() {
		AtomicInteger resourceCount = new AtomicInteger();
		withClient(createMcpTransport(), client -> {
			Flux<McpSchema.ReadResourceResult> resources = client.initialize()
				.then(client.listResources(null))
				.flatMapMany(r -> {
					List<Resource> l = r.resources();
					resourceCount.set(l.size());
					return Flux.fromIterable(l);
				})
				.flatMap(r -> client.readResource(r));

			StepVerifier.create(resources)
				.recordWith(ArrayList::new)
				.thenConsumeWhile(res -> true)
				.consumeRecordedWith(readResourceResults -> {
					assertThat(readResourceResults.size()).isEqualTo(resourceCount.get());
					for (ReadResourceResult result : readResourceResults) {

						assertThat(result).isNotNull();
						assertThat(result.contents()).isNotNull().isNotEmpty();

						// Validate each content item
						for (ResourceContents content : result.contents()) {
							assertThat(content).isNotNull();
							assertThat(content.uri()).isNotNull().isNotEmpty();
							assertThat(content.mimeType()).isNotNull().isNotEmpty();

							// Validate content based on its type with more comprehensive
							// checks
							switch (content.mimeType()) {
								case "text/plain" -> {
									TextResourceContents textContent = assertInstanceOf(TextResourceContents.class,
											content);
									assertThat(textContent.text()).isNotNull().isNotEmpty();
									assertThat(textContent.uri()).isNotEmpty();
								}
								case "application/octet-stream" -> {
									BlobResourceContents blobContent = assertInstanceOf(BlobResourceContents.class,
											content);
									assertThat(blobContent.blob()).isNotNull().isNotEmpty();
									assertThat(blobContent.uri()).isNotNull().isNotEmpty();
									// Validate base64 encoding format
									assertThat(blobContent.blob()).matches("^[A-Za-z0-9+/]*={0,2}$");
								}
								default -> {

									// Still validate basic properties
									if (content instanceof TextResourceContents textContent) {
										assertThat(textContent.text()).isNotNull();
									}
									else if (content instanceof BlobResourceContents blobContent) {
										assertThat(blobContent.blob()).isNotNull();
									}
								}
							}
						}
					}
				})
				.verifyComplete();
		});
	}

	@Test
	void testListResourceTemplatesWithoutInitialization() {
		verifyCallSucceedsWithImplicitInitialization(client -> client.listResourceTemplates(McpSchema.FIRST_PAGE),
				"listing resource templates");
	}

	@Test
	void testListResourceTemplates() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier
				.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listResourceTemplates(McpSchema.FIRST_PAGE)))
				.consumeNextWith(result -> {
					assertThat(result).isNotNull();
					assertThat(result.resourceTemplates()).isNotNull();
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllResourceTemplates() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listResourceTemplates()))
				.consumeNextWith(result -> {
					assertThat(result).isNotNull();
					assertThat(result.resourceTemplates()).isNotNull();
				})
				.verifyComplete();
		});
	}

	@Test
	void testListAllResourceTemplatesReturnsImmutableList() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.listResourceTemplates()))
				.consumeNextWith(result -> {
					assertThat(result.resourceTemplates()).isNotNull();
					// Verify that the returned list is immutable
					assertThatThrownBy(() -> result.resourceTemplates()
						.add(McpSchema.ResourceTemplate.builder("test://template", "test").title("test").build()))
						.isInstanceOf(UnsupportedOperationException.class);
				})
				.verifyComplete();
		});
	}

	@Test
	void testResourceSubscription() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.listResources().flatMap(resources -> {
				if (resources.resources().isEmpty()) {
					return Mono.empty();
				}
				Resource firstResource = resources.resources().get(0);
				return mcpAsyncClient.subscribeResource(SubscribeRequest.builder(firstResource.uri()).build())
					.then(mcpAsyncClient.unsubscribeResource(UnsubscribeRequest.builder(firstResource.uri()).build()));
			})).verifyComplete();
		});
	}

	@Test
	void testNotificationHandlers() {
		AtomicBoolean toolsNotificationReceived = new AtomicBoolean(false);
		AtomicBoolean resourcesNotificationReceived = new AtomicBoolean(false);
		AtomicBoolean promptsNotificationReceived = new AtomicBoolean(false);

		withClient(createMcpTransport(),
				builder -> builder
					.toolsChangeConsumer(tools -> Mono.fromRunnable(() -> toolsNotificationReceived.set(true)))
					.resourcesChangeConsumer(
							resources -> Mono.fromRunnable(() -> resourcesNotificationReceived.set(true)))
					.promptsChangeConsumer(prompts -> Mono.fromRunnable(() -> promptsNotificationReceived.set(true))),
				mcpAsyncClient -> {
					StepVerifier.create(mcpAsyncClient.initialize())
						.expectNextMatches(Objects::nonNull)
						.verifyComplete();
				});
	}

	@Test
	void testInitializeWithSamplingCapability() {
		ClientCapabilities capabilities = ClientCapabilities.builder().sampling().build();
		CreateMessageResult createMessageResult = CreateMessageResult
			.builder(McpSchema.Role.ASSISTANT, "test", "test-model")
			.build();
		withClient(createMcpTransport(),
				builder -> builder.capabilities(capabilities).sampling(request -> Mono.just(createMessageResult)),
				client -> {
					StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
				});
	}

	@Test
	void testInitializeWithElicitationCapability() {
		ClientCapabilities capabilities = ClientCapabilities.builder().elicitation().build();
		ElicitResult elicitResult = ElicitResult.builder(ElicitResult.Action.ACCEPT)
			.content(Map.of("foo", "bar"))
			.build();
		withClient(createMcpTransport(),
				builder -> builder.capabilities(capabilities).elicitation(request -> Mono.just(elicitResult)),
				client -> {
					StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
				});
	}

	@Test
	void testInitializeWithAllCapabilities() {
		var capabilities = ClientCapabilities.builder()
			.experimental(Map.of("feature", Map.of("featureFlag", true)))
			.roots(true)
			.sampling()
			.build();

		Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler = request -> Mono
			.just(CreateMessageResult.builder(McpSchema.Role.ASSISTANT, "test", "test-model").build());

		Function<ElicitFormRequest, Mono<ElicitResult>> formElicitationHandler = request -> Mono
			.just(ElicitResult.builder(ElicitResult.Action.ACCEPT).content(Map.of("foo", "bar")).build());

		withClient(createMcpTransport(),
				builder -> builder.capabilities(capabilities)
					.sampling(samplingHandler)
					.elicitation(formElicitationHandler),
				client ->

				StepVerifier.create(client.initialize()).assertNext(result -> {
					assertThat(result).isNotNull();
					assertThat(result.capabilities()).isNotNull();
				}).verifyComplete());
	}
	// ---------------------------------------
	// Logging Tests
	// ---------------------------------------

	@Test
	void testLoggingLevelsWithoutInitialization() {
		verifyNotificationSucceedsWithImplicitInitialization(
				client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG), "setting logging level");
	}

	@Test
	void testLoggingLevels() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier
				.create(mcpAsyncClient.initialize()
					.thenMany(Flux.fromArray(McpSchema.LoggingLevel.values()).flatMap(mcpAsyncClient::setLoggingLevel)))
				.verifyComplete();
		});
	}

	@Test
	void testLoggingConsumer() {
		AtomicBoolean logReceived = new AtomicBoolean(false);

		withClient(createMcpTransport(),
				builder -> builder.loggingConsumer(notification -> Mono.fromRunnable(() -> logReceived.set(true))),
				client -> {
					StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
				});

	}

	@Test
	void testLoggingWithNullNotification() {
		withClient(createMcpTransport(), mcpAsyncClient -> {
			StepVerifier.create(mcpAsyncClient.setLoggingLevel(null))
				.expectErrorMatches(error -> error.getMessage().contains("Logging level must not be null"))
				.verify();
		});
	}

	@Test
	void testSampling() {
		McpClientTransport transport = createMcpTransport();

		final String message = "Hello, world!";
		final String response = "Goodbye, world!";
		final int maxTokens = 100;

		AtomicReference<String> receivedPrompt = new AtomicReference<>();
		AtomicReference<String> receivedMessage = new AtomicReference<>();
		AtomicInteger receivedMaxTokens = new AtomicInteger();

		withClient(transport, spec -> spec.capabilities(McpSchema.ClientCapabilities.builder().sampling().build())
			.sampling(request -> {
				McpSchema.TextContent messageText = assertInstanceOf(McpSchema.TextContent.class,
						request.messages().get(0).content());
				receivedPrompt.set(request.systemPrompt());
				receivedMessage.set(messageText.text());
				receivedMaxTokens.set(request.maxTokens());

				return Mono.just(McpSchema.CreateMessageResult
					.builder(McpSchema.Role.USER, McpSchema.TextContent.builder(response).build(), "modelId")
					.stopReason(McpSchema.CreateMessageResult.StopReason.END_TURN)
					.build());
			}), client -> {
				StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();

				StepVerifier.create(client.callTool(McpSchema.CallToolRequest.builder("sampleLLM")
					.arguments(Map.of("prompt", message, "maxTokens", maxTokens))
					.build())).consumeNextWith(result -> {
						// Verify tool response to ensure our sampling response was passed
						// through
						assertThat(result.content()).hasAtLeastOneElementOfType(McpSchema.TextContent.class);
						assertThat(result.content()).allSatisfy(content -> {
							if (!(content instanceof McpSchema.TextContent text))
								return;

							assertThat(text.text()).contains(response);
						});

						// Verify sampling request parameters received in our callback
						assertThat(receivedPrompt.get()).isNotEmpty();
						assertThat(receivedMessage.get()).endsWith(message); // Prefixed
						assertThat(receivedMaxTokens.get()).isEqualTo(maxTokens);
					}).verifyComplete();
			});
	}

	// ---------------------------------------
	// Progress Notification Tests
	// ---------------------------------------

	@Test
	void testProgressConsumer() {
		Sinks.Many<McpSchema.ProgressNotification> sink = Sinks.many().unicast().onBackpressureBuffer();
		List<McpSchema.ProgressNotification> receivedNotifications = new CopyOnWriteArrayList<>();

		withClient(createMcpTransport(), builder -> builder.progressConsumer(notification -> {
			receivedNotifications.add(notification);
			sink.tryEmitNext(notification);
			return Mono.empty();
		}), client -> {
			StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();

			// Call a tool that sends progress notifications
			CallToolRequest request = CallToolRequest.builder()
				.name("longRunningOperation")
				.arguments(Map.of("duration", 1, "steps", 2))
				.progressToken("test-token")
				.build();

			StepVerifier.create(client.callTool(request)).consumeNextWith(result -> {
				assertThat(result).isNotNull();
			}).verifyComplete();

			// Use StepVerifier to verify the progress notifications via the sink
			StepVerifier.create(sink.asFlux()).expectNextCount(2).thenCancel().verify(Duration.ofSeconds(3));

			assertThat(receivedNotifications).hasSize(2);
			assertThat(receivedNotifications.get(0).progressToken()).isEqualTo("test-token");
		});
	}

}


//! xAI Completion Integration
//!
//! Uses the xAI Responses API: <https://docs.x.ai/docs/guides/chat>

use crate::telemetry::{CompletionOperation, CompletionSpanBuilder, SpanCombinator};
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tracing::{Instrument, Level, enabled};

use super::api::{ApiResponse, Message, ToolDefinition};
use super::client::Client;
use crate::OneOrMany;
use crate::completion::{self, CompletionError, CompletionRequest, GetTokenUsage};
use crate::http_client::HttpClientExt;
use crate::providers::openai::responses_api::ToolChoice;
use crate::providers::openai::responses_api::streaming::StreamingCompletionResponse;
use crate::providers::openai::responses_api::{Output, ResponsesUsage};
use crate::streaming::StreamingCompletionResponse as BaseStreamingCompletionResponse;

/// xAI completion models as of 2025-06-04
pub const GROK_2_1212: &str = "grok-2-1212";
pub const GROK_2_VISION_1212: &str = "grok-2-vision-1212";
pub const GROK_3: &str = "grok-3";
pub const GROK_3_FAST: &str = "grok-3-fast";
pub const GROK_3_MINI: &str = "grok-3-mini";
pub const GROK_3_MINI_FAST: &str = "grok-3-mini-fast";
pub const GROK_2_IMAGE_1212: &str = "grok-2-image-1212";
pub const GROK_4: &str = "grok-4-0709";

// ================================================================
// Request Types
// ================================================================

#[derive(Debug, Serialize, Deserialize)]
pub(super) struct XAICompletionRequest {
    pub(super) model: String,
    pub input: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_output_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_choice: Option<ToolChoice>,
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub additional_params: Option<serde_json::Value>,
}

impl TryFrom<(&str, CompletionRequest)> for XAICompletionRequest {
    type Error = CompletionError;

    fn try_from((model, req): (&str, CompletionRequest)) -> Result<Self, Self::Error> {
        let chat_history = req.chat_history_with_documents();
        if req.output_schema.is_some() {
            tracing::warn!("Structured outputs currently not supported for xAI");
        }
        let model = req.model.clone().unwrap_or_else(|| model.to_string());
        let mut input: Vec<Message> = req
            .preamble
            .as_ref()
            .map_or_else(Vec::new, |p| vec![Message::system(p)]);

        let mut additional_params_payload = req.additional_params.unwrap_or(Value::Null);

        for msg in chat_history {
            let msg: Vec<Message> = msg.try_into()?;
            input.extend(msg);
        }

        let tool_choice = req.tool_choice.map(ToolChoice::try_from).transpose()?;
        let mut additional_tools =
            extract_tools_from_additional_params(&mut additional_params_payload)?;
        let mut tools = req
            .tools
            .into_iter()
            .map(ToolDefinition::from)
            .map(serde_json::to_value)
            .collect::<Result<Vec<_>, _>>()?;
        tools.append(&mut additional_tools);
        let additional_params = if additional_params_payload.is_null() {
            None
        } else {
            Some(additional_params_payload)
        };

        Ok(Self {
            model: model.to_string(),
            input,
            temperature: req.temperature,
            max_output_tokens: req.max_tokens,
            tools,
            tool_choice,
            additional_params,
        })
    }
}

fn extract_tools_from_additional_params(
    additional_params: &mut Value,
) -> Result<Vec<Value>, CompletionError> {
    if let Some(map) = additional_params.as_object_mut()
        && let Some(raw_tools) = map.remove("tools")
    {
        return serde_json::from_value::<Vec<Value>>(raw_tools).map_err(|err| {
            CompletionError::RequestError(
                format!("Invalid xAI `additional_params.tools` payload: {err}").into(),
            )
        });
    }

    Ok(Vec::new())
}

// ================================================================
// Response Types
// ================================================================

#[derive(Debug, Deserialize, Serialize)]
pub struct CompletionResponse {
    pub id: String,
    pub model: String,
    pub output: Vec<Output>,
    #[serde(default)]
    pub created: i64,
    #[serde(default)]
    pub object: String,
    #[serde(default)]
    pub status: Option<String>,
    pub usage: Option<ResponsesUsage>,
}

impl TryFrom<CompletionResponse> for completion::CompletionResponse<CompletionResponse> {
    type Error = CompletionError;

    fn try_from(response: CompletionResponse) -> Result<Self, Self::Error> {
        let content: Vec<completion::AssistantContent> = response
            .output
            .iter()
            .cloned()
            .flat_map(<Vec<completion::AssistantContent>>::from)
            .collect();

        let choice = OneOrMany::many(content).map_err(|_| {
            CompletionError::ResponseError("Response contained no output".to_owned())
        })?;

        let usage = response
            .usage
            .as_ref()
            .map(GetTokenUsage::token_usage)
            .unwrap_or_default();
        let message_id = response.output.iter().find_map(|item| match item {
            Output::Message(message) => Some(message.id.clone()),
            _ => None,
        });

        Ok(completion::CompletionResponse {
            choice,
            usage,
            raw_response: response,
            message_id,
        })
    }
}

// ================================================================
// Completion Model
// ================================================================

#[derive(Clone)]
pub struct CompletionModel<T = reqwest::Client> {
    pub(crate) client: Client<T>,
    pub model: String,
}

impl<T> CompletionModel<T> {
    pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
        Self {
            client,
            model: model.into(),
        }
    }
}

impl<T> completion::CompletionModel for CompletionModel<T>
where
    T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
{
    type Response = CompletionResponse;
    type StreamingResponse = StreamingCompletionResponse;

    type Client = Client<T>;

    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
        Self::new(client.clone(), model)
    }

    async fn completion(
        &self,
        completion_request: completion::CompletionRequest,
    ) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
        let system_instructions = completion_request.preamble.clone();
        let record_telemetry_content = completion_request.record_telemetry_content;
        let request =
            XAICompletionRequest::try_from((self.model.to_string().as_ref(), completion_request))?;
        let span = CompletionSpanBuilder::new("xai", &request.model, CompletionOperation::Chat)
            .system_instructions(system_instructions.as_deref(), record_telemetry_content)
            .build();

        if enabled!(Level::TRACE) {
            tracing::trace!(target: "rig::completions",
                "xAI completion request: {}",
                serde_json::to_string_pretty(&request)?
            );
        }

        let body = serde_json::to_vec(&request)?;
        let req = self
            .client
            .post("/v1/responses")?
            .body(body)
            .map_err(|e| CompletionError::HttpError(e.into()))?;

        async move {
            let response = self.client.send::<_, Bytes>(req).await?;
            let status = response.status();
            let response_body = response.into_body().into_future().await?.to_vec();

            if status.is_success() {
                match serde_json::from_slice::<ApiResponse<CompletionResponse>>(&response_body)? {
                    ApiResponse::Ok(response) => {
                        let span = tracing::Span::current();
                        span.record("gen_ai.response.id", response.id.as_str());
                        span.record("gen_ai.response.model", response.model.as_str());
                        if let Some(usage) = &response.usage {
                            span.record_token_usage(usage);
                        }

                        if enabled!(Level::TRACE) {
                            tracing::trace!(target: "rig::completions",
                                "xAI completion response: {}",
                                serde_json::to_string_pretty(&response)?
                            );
                        }

                        response.try_into()
                    }
                    ApiResponse::Error(error) => {
                        tracing::warn!(message = %error.message(), "provider returned an error response");
                        Err(CompletionError::from_http_response(
                            status,
                            String::from_utf8_lossy(&response_body),
                        ))
                    }
                }
            } else {
                Err(CompletionError::from_http_response(
                    status,
                    String::from_utf8_lossy(&response_body),
                ))
            }
        }
        .instrument(span)
        .await
    }

    async fn stream(
        &self,
        request: CompletionRequest,
    ) -> Result<BaseStreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
        self.stream(request).await
    }
}

#[cfg(test)]
mod tests {
    use super::XAICompletionRequest;
    use crate::OneOrMany;
    use crate::completion::request::Document;
    use crate::completion::{CompletionRequest, CompletionRequestBuilder, Message, ToolDefinition};
    use crate::message::ToolChoice;
    use crate::test_utils::MockCompletionModel;

    #[test]
    fn xai_request_includes_normalized_documents() {
        let request =
            CompletionRequestBuilder::new(MockCompletionModel::default(), "What is glarb-glarb?")
                .message(Message::system("Use the provided context."))
                .document(Document {
                    id: "doc_1".to_string(),
                    text: "Definition of glarb-glarb: an ancient tool.".to_string(),
                    additional_props: Default::default(),
                })
                .build();

        let xai_request = XAICompletionRequest::try_from(("grok-4-0709", request))
            .expect("request conversion should succeed");
        let serialized = serde_json::to_value(xai_request).expect("serialization should succeed");
        let input = serialized["input"]
            .as_array()
            .expect("xAI request input should be an array");

        assert!(
            input
                .iter()
                .any(|message| message.to_string().contains("glarb-glarb")),
            "normalized documents should be forwarded into xAI input"
        );
    }

    #[test]
    fn xai_direct_request_keeps_documents_after_system_messages() {
        let request = CompletionRequest {
            model: None,
            preamble: None,
            chat_history: OneOrMany::many(vec![
                Message::system("System prompt"),
                Message::assistant("Earlier assistant turn"),
                Message::system("Mid-conversation instruction"),
                Message::user("What is glarb-glarb?"),
            ])
            .unwrap(),
            documents: vec![Document {
                id: "doc_1".to_string(),
                text: "Definition of glarb-glarb: an ancient tool.".to_string(),
                additional_props: Default::default(),
            }],
            tools: vec![],
            temperature: None,
            max_tokens: None,
            tool_choice: None,
            additional_params: None,
            output_schema: None,
            record_telemetry_content: false,
        };

        let xai_request = XAICompletionRequest::try_from(("grok-4-0709", request))
            .expect("request conversion should succeed");
        let serialized = serde_json::to_value(xai_request).expect("serialization should succeed");
        let input = serialized["input"]
            .as_array()
            .expect("xAI request input should be an array");

        assert_eq!(input.len(), 5);
        assert_eq!(input[0]["role"], "system");
        assert_eq!(input[1]["role"], "user");
        assert!(
            input[1].to_string().contains("<file id: doc_1>"),
            "document input should follow leading system input: {input:?}"
        );
        assert_eq!(input[2]["role"], "assistant");
        assert_eq!(input[3]["role"], "system");
        assert_eq!(input[4]["role"], "user");
        assert_eq!(
            input
                .iter()
                .filter(|message| message.to_string().contains("<file id: doc_1>"))
                .count(),
            1,
            "document input should appear exactly once: {input:?}"
        );
    }

    #[test]
    fn xai_request_uses_responses_tool_choice_for_specific_tool() {
        let request = CompletionRequestBuilder::new(MockCompletionModel::default(), "Use a tool.")
            .tool(ToolDefinition {
                name: "alpha".to_string(),
                description: "Alpha tool".to_string(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": []
                }),
            })
            .tool(ToolDefinition {
                name: "beta".to_string(),
                description: "Beta tool".to_string(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": []
                }),
            })
            .tool_choice(ToolChoice::Specific {
                function_names: vec!["beta".to_string()],
            })
            .build();

        let xai_request = XAICompletionRequest::try_from(("grok-4.3", request))
            .expect("xAI Responses API should support specific tool choice");
        let serialized = serde_json::to_value(xai_request).expect("serialization should succeed");

        assert_eq!(
            serialized["tool_choice"],
            serde_json::json!({"type": "function", "name": "beta"})
        );
    }

    #[test]
    fn xai_response_preserves_message_id_and_reasoning_token_usage() {
        let raw: super::CompletionResponse = serde_json::from_value(serde_json::json!({
            "id": "resp_123",
            "model": "grok-4.3",
            "output": [
                {
                    "type": "reasoning",
                    "id": "rs_123",
                    "summary": [{ "type": "summary_text", "text": "thinking" }],
                    "status": "completed"
                },
                {
                    "type": "message",
                    "id": "msg_123",
                    "role": "assistant",
                    "status": "completed",
                    "content": [
                        { "type": "output_text", "text": "done", "annotations": [] }
                    ]
                }
            ],
            "usage": {
                "input_tokens": 10,
                "input_tokens_details": { "cached_tokens": 3 },
                "output_tokens": 8,
                "output_tokens_details": { "reasoning_tokens": 5 },
                "total_tokens": 18
            }
        }))
        .expect("fixture should deserialize");

        let converted = crate::completion::CompletionResponse::try_from(raw)
            .expect("xAI response should convert");

        assert_eq!(converted.message_id.as_deref(), Some("msg_123"));
        assert_eq!(converted.usage.input_tokens, 10);
        assert_eq!(converted.usage.cached_input_tokens, 3);
        assert_eq!(converted.usage.output_tokens, 8);
        assert_eq!(converted.usage.reasoning_tokens, 5);
    }

    #[tokio::test]
    async fn completion_non_success_preserves_status_and_body() {
        use crate::client::CompletionClient;
        use crate::completion::{CompletionError, CompletionModel as _};
        use crate::test_utils::RecordingHttpClient;

        let body = r#"{"error":"boom","code":"503"}"#;
        let http_client =
            RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
        let client = crate::providers::xai::Client::builder()
            .api_key("test-key")
            .http_client(http_client)
            .build()
            .expect("build client");
        let model = client.completion_model(crate::providers::xai::completion::GROK_4);
        let request = model.completion_request("hello").build();

        let error = model
            .completion(request)
            .await
            .expect_err("should fail with non-success status");

        assert!(matches!(error, CompletionError::HttpError(_)));
        assert_eq!(
            error.provider_response_status(),
            Some(http::StatusCode::SERVICE_UNAVAILABLE)
        );
        assert_eq!(error.provider_response_body(), Some(body));
    }

    #[tokio::test]
    async fn completion_2xx_error_envelope_preserves_status_and_body() {
        use crate::client::CompletionClient;
        use crate::completion::{CompletionError, CompletionModel as _};
        use crate::test_utils::RecordingHttpClient;

        // Deserializes to `ApiResponse::Error(ApiError { error, code })` on a 200 OK.
        let body = r#"{"error":"boom","code":"503"}"#;
        let http_client = RecordingHttpClient::new(body);
        let client = crate::providers::xai::Client::builder()
            .api_key("test-key")
            .http_client(http_client)
            .build()
            .expect("build client");
        let model = client.completion_model(crate::providers::xai::completion::GROK_4);
        let request = model.completion_request("hello").build();

        let error = model
            .completion(request)
            .await
            .expect_err("should fail with provider error envelope");

        match &error {
            CompletionError::ProviderResponse(stored) => {
                assert_eq!(stored.body, body);
                assert_eq!(stored.status, Some(http::StatusCode::OK));
            }
            other => panic!("expected ProviderResponse, got {other:?}"),
        }
    }
}


//go:build linux

package agent

import (
	"os"
	"path/filepath"
	"testing"

	"github.com/henrygd/beszel/internal/entities/smart"
)

func TestEmmcMockSysfsScanAndCollect(t *testing.T) {
	tmp := t.TempDir()
	prev := emmcSysfsRoot
	emmcSysfsRoot = tmp
	t.Cleanup(func() { emmcSysfsRoot = prev })

	// Fake: /sys/class/block/mmcblk0
	mmcDeviceDir := filepath.Join(tmp, "class", "block", "mmcblk0", "device")
	mmcQueueDir := filepath.Join(tmp, "class", "block", "mmcblk0", "queue")
	if err := os.MkdirAll(mmcDeviceDir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.MkdirAll(mmcQueueDir, 0o755); err != nil {
		t.Fatal(err)
	}

	write := func(path, content string) {
		t.Helper()
		if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
			t.Fatal(err)
		}
	}

	write(filepath.Join(mmcDeviceDir, "pre_eol_info"), "0x02\n")
	write(filepath.Join(mmcDeviceDir, "life_time"), "0x04 0x05\n")
	write(filepath.Join(mmcDeviceDir, "name"), "H26M52103FMR\n")
	write(filepath.Join(mmcDeviceDir, "serial"), "01234567\n")
	write(filepath.Join(mmcDeviceDir, "prv"), "0x08\n")
	write(filepath.Join(mmcQueueDir, "logical_block_size"), "512\n")
	write(filepath.Join(tmp, "class", "block", "mmcblk0", "size"), "1024\n") // sectors

	devs := scanEmmcDevices()
	if len(devs) != 1 {
		t.Fatalf("scanEmmcDevices() = %d devices, want 1", len(devs))
	}
	if devs[0].Name != "/dev/mmcblk0" || devs[0].Type != "emmc" {
		t.Fatalf("scanEmmcDevices()[0] = %+v, want Name=/dev/mmcblk0 Type=emmc", devs[0])
	}

	sm := &SmartManager{SmartDataMap: map[string]*smart.SmartData{}}
	ok, err := sm.collectEmmcHealth(devs[0])
	if err != nil || !ok {
		t.Fatalf("collectEmmcHealth() = (ok=%v, err=%v), want (true,nil)", ok, err)
	}
	if len(sm.SmartDataMap) != 1 {
		t.Fatalf("SmartDataMap len=%d, want 1", len(sm.SmartDataMap))
	}
	var got *smart.SmartData
	for _, v := range sm.SmartDataMap {
		got = v
		break
	}
	if got == nil {
		t.Fatalf("SmartDataMap value nil")
	}
	if got.DiskType != "emmc" || got.DiskName != "/dev/mmcblk0" {
		t.Fatalf("disk fields = (type=%q name=%q), want (emmc,/dev/mmcblk0)", got.DiskType, got.DiskName)
	}
	if got.SmartStatus != "WARNING" {
		t.Fatalf("SmartStatus=%q, want WARNING", got.SmartStatus)
	}
	if got.SerialNumber != "01234567" || got.ModelName == "" || got.Capacity == 0 {
		t.Fatalf("identity fields = (model=%q serial=%q cap=%d), want non-empty model, serial 01234567, cap>0", got.ModelName, got.SerialNumber, got.Capacity)
	}
	if len(got.Attributes) < 3 {
		t.Fatalf("attributes len=%d, want >= 3", len(got.Attributes))
	}
}


/*
 * Copyright 2024 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package compose

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"
	"testing"
	"unicode/utf8"

	"github.com/stretchr/testify/assert"

	"github.com/cloudwego/eino/components/prompt"
	"github.com/cloudwego/eino/schema"
)

func TestChainBranch(t *testing.T) {
	cond := func(ctx context.Context, input string) (key string, err error) {
		switch input {
		case "one":
			return "one_key", nil
		case "two":
			return "two_key", nil
		case "three":
			return "three_key", nil
		default:
			return "", fmt.Errorf("invalid input= %s", input)
		}
	}

	t.Run("nested chain", func(t *testing.T) {
		inner := NewChain[string, string]()
		inner.AppendBranch(NewChainBranch(cond).
			AddLambda("one_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
				return in + in, nil
			})).
			AddLambda("two_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
				return in + in + in, nil
			})))
		inner.AppendParallel(NewParallel().
			AddLambda("one_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
				return in + in, nil
			})).
			AddLambda("two_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
				return in + in + in, nil
			})))

		outer := NewChain[string, string]()
		outer.AppendGraph(inner)
		_, err := outer.Compile(context.Background())
		assert.Error(t, err)
	})

	t.Run("bad param", func(t *testing.T) {
		c := NewChain[string, string]()
		c.AppendBranch(nil)
		assert.NotNil(t, c.err)

		c = NewChain[string, string]()
		c.AppendBranch(NewChainBranch[string](nil))
		assert.NotNil(t, c.err)

		c = NewChain[string, string]()
		c.AppendBranch(NewChainBranch(cond).AddChatTemplate("template", prompt.FromMessages(schema.FString, schema.SystemMessage("hello"))))
		assert.NotNil(t, c.err)

		c = NewChain[string, string]()
		c.AppendBranch(NewChainBranch(cond).AddChatTemplate("1", prompt.FromMessages(schema.FString)).AddChatTemplate("1", prompt.FromMessages(schema.FString)))
		assert.NotNil(t, c.err)
	})

	t.Run("different Node types in branch", func(t *testing.T) {
		c := NewChain[string, string]()
		c.AppendBranch(NewChainBranch(cond).
			AddChatTemplate("t", prompt.FromMessages(schema.FString)).
			AddGraph("c", NewChain[string, string]()))
		assert.NotNil(t, c.err)
	})

	t.Run("type mismatch", func(t *testing.T) {
		c := NewChain[int, string]()
		c.AppendBranch(NewChainBranch(cond).
			AddLambda("one_key", InvokableLambda(func(ctx context.Context, in int) (output string, err error) {
				return strconv.Itoa(in), nil
			})).
			AddLambda("two_key", InvokableLambda(func(ctx context.Context, in int) (output string, err error) {
				return strconv.Itoa(in), nil
			})))
		_, err := c.Compile(context.Background())
		assert.NotNil(t, err)
	})

	t.Run("invoke", func(t *testing.T) {
		c := NewChain[string, string]()
		c.AppendBranch(NewChainBranch(cond).
			AddLambda("one_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
				return in + in, nil
			})).
			AddLambda("two_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
				return in + in + in, nil
			})))
		c.AppendLambda(InvokableLambda(func(ctx context.Context, in string) (output string, err error) {
			return in + in, nil
		}))
		assert.Nil(t, c.err)
		compiledChain, err := c.Compile(context.Background())
		assert.Nil(t, err)

		out, err := compiledChain.Invoke(context.Background(), "two")
		assert.Nil(t, err)
		assert.Equal(t, "twotwotwotwotwotwo", out)

		_, err = compiledChain.Invoke(context.Background(), "three")
		assert.NotNil(t, err)

		_, err = compiledChain.Invoke(context.Background(), "four")
		assert.NotNil(t, err)
	})

	t.Run("fake stream", func(t *testing.T) {
		c := NewChain[string, string]()
		c.AppendLambda(StreamableLambda(func(ctx context.Context, in string) (output *schema.StreamReader[string], err error) {
			sr, sw := schema.Pipe[string](utf8.RuneCountInString(in))

			go func() {
				for _, field := range strings.Fields(in) {
					sw.Send(field, nil)
				}
				sw.Close()
			}()

			return sr, nil
		}))
		c.AppendBranch(NewChainBranch[string](cond).AddLambda("one_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {
			defer in.Close()
			for {
				v, err := in.Recv()
				if errors.Is(err, io.EOF) {
					break
				}

				if err != nil {
					return "", err
				}

				output += v
			}

			return output + output, nil
		})).
			AddLambda("two_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {
				defer in.Close()
				for {
					v, err := in.Recv()
					if errors.Is(err, io.EOF) {
						break
					}

					if err != nil {
						return "", err
					}

					output += v
				}

				return output + output + output, nil
			})))

		assert.Nil(t, c.err)
		compiledChain, err := c.Compile(context.Background())
		assert.Nil(t, err)

		out, err := compiledChain.Invoke(context.Background(), "one")
		assert.Nil(t, err)
		assert.Equal(t, "oneone", out)
	})

	t.Run("real stream", func(t *testing.T) {
		streamCon := func(ctx context.Context, sr *schema.StreamReader[string]) (key string, err error) {
			msg, err := sr.Recv()
			if err != nil {
				return "", err
			}
			defer sr.Close()

			switch msg {
			case "one":
				return "one_key", nil
			case "two":
				return "two_key", nil
			case "three":
				return "three_key", nil
			default:
				return "", fmt.Errorf("invalid input= %s", msg)
			}
		}

		c := NewChain[string, string]()
		c.AppendLambda(StreamableLambda(func(ctx context.Context, in string) (output *schema.StreamReader[string], err error) {
			sr, sw := schema.Pipe[string](utf8.RuneCountInString(in))

			go func() {
				for _, field := range strings.Fields(in) {
					sw.Send(field, nil)
				}
				sw.Close()
			}()

			return sr, nil
		}))
		c.AppendBranch(NewStreamChainBranch(streamCon).AddLambda("one_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {
			defer in.Close()
			for {
				v, err := in.Recv()
				if errors.Is(err, io.EOF) {
					break
				}

				if err != nil {
					return "", err
				}

				output += v
			}

			return output + output, nil
		})).
			AddLambda("two_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) {
				defer in.Close()
				for {
					v, err := in.Recv()
					if errors.Is(err, io.EOF) {
						break
					}

					if err != nil {
						return "", err
					}

					output += v
				}

				return output + output + output, nil
			})))

		assert.Nil(t, c.err)
		compiledChain, err := c.Compile(context.Background())
		assert.Nil(t, err)

		out, err := compiledChain.Stream(context.Background(), "one size fit all")
		assert.Nil(t, err)
		concat, err := concatStreamReader(out)
		assert.Nil(t, err)
		assert.Equal(t, "onesizefitallonesizefitall", concat)
	})
}

func TestChainMultiBranch(t *testing.T) {
	emptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil })

	ctx := context.Background()
	r, err := NewChain[string, map[string]any]().
		AppendBranch(NewChainMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) {
			return map[string]bool{"1": true, "2": true}, nil
		}).AddLambda("1", emptyLambda, WithOutputKey("1")).AddLambda("2", emptyLambda, WithOutputKey("2")).AddLambda("3", emptyLambda, WithOutputKey("3"))).
		Compile(ctx)
	assert.Nil(t, err)

	result, err := r.Invoke(ctx, "start")
	assert.NoError(t, err)
	assert.Equal(t, map[string]any{
		"1": "start",
		"2": "start",
	}, result)

	streamResult, err := r.Stream(ctx, "start")
	assert.NoError(t, err)
	result = map[string]any{}
	for {
		chunk, err := streamResult.Recv()
		if err == io.EOF {
			break
		}
		assert.NoError(t, err)
		for k, v := range chunk {
			result[k] = v
		}
	}
	assert.Equal(t, map[string]any{
		"1": "start",
		"2": "start",
	}, result)
}

func TestStreamChainMultiBranch(t *testing.T) {
	emptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil })

	ctx := context.Background()
	r, err := NewChain[string, map[string]any]().
		AppendBranch(NewStreamChainMultiBranch(func(ctx context.Context, in *schema.StreamReader[string]) (endNode map[string]bool, err error) {
			return map[string]bool{"1": true, "2": true}, nil
		}).AddLambda("1", emptyLambda, WithOutputKey("1")).AddLambda("2", emptyLambda, WithOutputKey("2")).AddLambda("3", emptyLambda, WithOutputKey("3"))).
		Compile(ctx)
	assert.Nil(t, err)

	result, err := r.Invoke(ctx, "start")
	assert.NoError(t, err)
	assert.Equal(t, map[string]any{
		"1": "start",
		"2": "start",
	}, result)

	streamResult, err := r.Stream(ctx, "start")
	assert.NoError(t, err)
	result = map[string]any{}
	for {
		chunk, err := streamResult.Recv()
		if err == io.EOF {
			break
		}
		assert.NoError(t, err)
		for k, v := range chunk {
			result[k] = v
		}
	}
	assert.Equal(t, map[string]any{
		"1": "start",
		"2": "start",
	}, result)
}


import sys
import re
import os
from typing import BinaryIO, Any, List
from enum import Enum

from .._base_converter import DocumentConverter, DocumentConverterResult
from .._stream_info import StreamInfo
from .._exceptions import MissingDependencyException

# Try loading optional (but in this case, required) dependencies
# Save reporting of any exceptions for later
_dependency_exc_info = None
try:
    from azure.ai.documentintelligence import DocumentIntelligenceClient
    from azure.ai.documentintelligence.models import (
        AnalyzeDocumentRequest,
        AnalyzeResult,
        DocumentAnalysisFeature,
    )
    from azure.core.credentials import AzureKeyCredential, TokenCredential
    from azure.identity import DefaultAzureCredential
except ImportError:
    # Preserve the error and stack trace for later
    _dependency_exc_info = sys.exc_info()

    # Define these types for type hinting when the package is not available
    class AzureKeyCredential:
        pass

    class TokenCredential:
        pass

    class DocumentIntelligenceClient:
        pass

    class AnalyzeDocumentRequest:
        pass

    class AnalyzeResult:
        pass

    class DocumentAnalysisFeature:
        pass

    class DefaultAzureCredential:
        pass

# TODO: currently, there is a bug in the document intelligence SDK with importing the "ContentFormat" enum.
# This constant is a temporary fix until the bug is resolved.
CONTENT_FORMAT = "markdown"

class DocumentIntelligenceFileType(str, Enum):
    """Enum of file types supported by the Document Intelligence Converter."""

    # No OCR
    DOCX = "docx"
    PPTX = "pptx"
    XLSX = "xlsx"
    HTML = "html"
    # OCR
    PDF = "pdf"
    JPEG = "jpeg"
    PNG = "png"
    BMP = "bmp"
    TIFF = "tiff"

def _get_mime_type_prefixes(types: List[DocumentIntelligenceFileType]) -> List[str]:
    """Get the MIME type prefixes for the given file types."""
    prefixes: List[str] = []
    for type_ in types:
        if type_ == DocumentIntelligenceFileType.DOCX:
            prefixes.append(
                "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            )
        elif type_ == DocumentIntelligenceFileType.PPTX:
            prefixes.append(
                "application/vnd.openxmlformats-officedocument.presentationml"
            )
        elif type_ == DocumentIntelligenceFileType.XLSX:
            prefixes.append(
                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            )
        elif type_ == DocumentIntelligenceFileType.HTML:
            prefixes.append("text/html")
            prefixes.append("application/xhtml+xml")
        elif type_ == DocumentIntelligenceFileType.PDF:
            prefixes.append("application/pdf")
            prefixes.append("application/x-pdf")
        elif type_ == DocumentIntelligenceFileType.JPEG:
            prefixes.append("image/jpeg")
        elif type_ == DocumentIntelligenceFileType.PNG:
            prefixes.append("image/png")
        elif type_ == DocumentIntelligenceFileType.BMP:
            prefixes.append("image/bmp")
        elif type_ == DocumentIntelligenceFileType.TIFF:
            prefixes.append("image/tiff")
    return prefixes

def _get_file_extensions(types: List[DocumentIntelligenceFileType]) -> List[str]:
    """Get the file extensions for the given file types."""
    extensions: List[str] = []
    for type_ in types:
        if type_ == DocumentIntelligenceFileType.DOCX:
            extensions.append(".docx")
        elif type_ == DocumentIntelligenceFileType.PPTX:
            extensions.append(".pptx")
        elif type_ == DocumentIntelligenceFileType.XLSX:
            extensions.append(".xlsx")
        elif type_ == DocumentIntelligenceFileType.PDF:
            extensions.append(".pdf")
        elif type_ == DocumentIntelligenceFileType.JPEG:
            extensions.append(".jpg")
            extensions.append(".jpeg")
        elif type_ == DocumentIntelligenceFileType.PNG:
            extensions.append(".png")
        elif type_ == DocumentIntelligenceFileType.BMP:
            extensions.append(".bmp")
        elif type_ == DocumentIntelligenceFileType.TIFF:
            extensions.append(".tiff")
        elif type_ == DocumentIntelligenceFileType.HTML:
            extensions.append(".html")
    return extensions

class DocumentIntelligenceConverter(DocumentConverter):
    """Specialized DocumentConverter that uses Document Intelligence to extract text from documents."""

    def __init__(
        self,
        *,
        endpoint: str,
        api_version: str = "2024-07-31-preview",
        credential: AzureKeyCredential | TokenCredential | None = None,
        file_types: List[DocumentIntelligenceFileType] = [
            DocumentIntelligenceFileType.DOCX,
            DocumentIntelligenceFileType.PPTX,
            DocumentIntelligenceFileType.XLSX,
            DocumentIntelligenceFileType.PDF,
            DocumentIntelligenceFileType.JPEG,
            DocumentIntelligenceFileType.PNG,
            DocumentIntelligenceFileType.BMP,
            DocumentIntelligenceFileType.TIFF,
        ],
    ):
        """
        Initialize the DocumentIntelligenceConverter.

        Args:
            endpoint (str): The endpoint for the Document Intelligence service.
            api_version (str): The API version to use. Defaults to "2024-07-31-preview".
            credential (AzureKeyCredential | TokenCredential | None): The credential to use for authentication.
            file_types (List[DocumentIntelligenceFileType]): The file types to accept. Defaults to all supported file types.
        """

        super().__init__()
        self._file_types = file_types

        # Raise an error if the dependencies are not available.
        # This is different than other converters since this one isn't even instantiated
        # unless explicitly requested.
        if _dependency_exc_info is not None:
            raise MissingDependencyException(
                "DocumentIntelligenceConverter requires the optional dependency [az-doc-intel] (or [all]) to be installed. E.g., `pip install 'markitdown[az-doc-intel]'`"
            ) from _dependency_exc_info[
                1
            ].with_traceback(  # type: ignore[union-attr]
                _dependency_exc_info[2]
            )

        if credential is None:
            if os.environ.get("AZURE_API_KEY") is None:
                credential = DefaultAzureCredential()
            else:
                credential = AzureKeyCredential(os.environ["AZURE_API_KEY"])

        self.endpoint = endpoint
        self.api_version = api_version
        self.doc_intel_client = DocumentIntelligenceClient(
            endpoint=self.endpoint,
            api_version=self.api_version,
            credential=credential,
        )

    def accepts(
        self,
        file_stream: BinaryIO,
        stream_info: StreamInfo,
        **kwargs: Any,  # Options to pass to the converter
    ) -> bool:
        mimetype = (stream_info.mimetype or "").lower()
        extension = (stream_info.extension or "").lower()

        if extension in _get_file_extensions(self._file_types):
            return True

        for prefix in _get_mime_type_prefixes(self._file_types):
            if mimetype.startswith(prefix):
                return True

        return False

    def _analysis_features(self, stream_info: StreamInfo) -> List[str]:
        """
        Helper needed to determine which analysis features to use.
        Certain document analysis features are not available for
        office filetypes (.xlsx, .pptx, .html, .docx)
        """
        mimetype = (stream_info.mimetype or "").lower()
        extension = (stream_info.extension or "").lower()

        # Types that don't support ocr
        no_ocr_types = [
            DocumentIntelligenceFileType.DOCX,
            DocumentIntelligenceFileType.PPTX,
            DocumentIntelligenceFileType.XLSX,
            DocumentIntelligenceFileType.HTML,
        ]

        if extension in _get_file_extensions(no_ocr_types):
            return []

        for prefix in _get_mime_type_prefixes(no_ocr_types):
            if mimetype.startswith(prefix):
                return []

        return [
            DocumentAnalysisFeature.FORMULAS,  # enable formula extraction
            DocumentAnalysisFeature.OCR_HIGH_RESOLUTION,  # enable high resolution OCR
            DocumentAnalysisFeature.STYLE_FONT,  # enable font style extraction
        ]

    def convert(
        self,
        file_stream: BinaryIO,
        stream_info: StreamInfo,
        **kwargs: Any,  # Options to pass to the converter
    ) -> DocumentConverterResult:
        # Extract the text using Azure Document Intelligence
        poller = self.doc_intel_client.begin_analyze_document(
            model_id="prebuilt-layout",
            body=AnalyzeDocumentRequest(bytes_source=file_stream.read()),
            features=self._analysis_features(stream_info),
            output_content_format=CONTENT_FORMAT,  # TODO: replace with "ContentFormat.MARKDOWN" when the bug is fixed
        )
        result: AnalyzeResult = poller.result()

        # remove comments from the markdown content generated by Doc Intelligence and append to markdown string
        markdown_text = re.sub(r"<!--.*?-->", "", result.content, flags=re.DOTALL)
        return DocumentConverterResult(markdown=markdown_text)


// caveman-shrink — pure-Node prose compressor for MCP tool descriptions
// and other safe string fields. Mirrors the boundaries of the
// caveman-compress Python tool (preserve code, URLs, paths, identifiers)
// but reimplemented in Node so the proxy stays single-runtime.
//
// API: compress(text, opts?) → { compressed, before, after }
//
// Boundaries (NEVER touched):
//  - fenced code blocks (``` ... ```)
//  - inline code (`...`)
//  - URLs (https?://...)
//  - filesystem paths (anything with `/` or `\`)
//  - "code-looking" tokens (parens at end, semicolons, JSON-like braces)
//  - identifiers in CamelCase / snake_case / dotted.path form
//
// Compression applied to everything else:
//  - drop articles (a, an, the)
//  - drop filler words (just, really, basically, actually, simply, quite, very)
//  - drop pleasantries (please, kindly, thank you, sure, certainly)
//  - drop hedging (perhaps, maybe, might, could potentially, would like to)
//  - drop leading "I'll" / "I will" / "you can" / "we will" / "let me"
//  - collapse whitespace runs

const FILLERS = new RegExp(
  '\\b(?:just|really|basically|actually|simply|quite|very|essentially|literally)\\b',
  'gi'
);

const PLEASANTRIES = new RegExp(
  '\\b(?:please|kindly|thank you|thanks|sure|certainly|of course|happy to|i\'?d be happy)\\b[,.]?\\s*',
  'gi'
);

const HEDGES = new RegExp(
  '\\b(?:perhaps|maybe|might|could potentially|would like to|i think|in my opinion|it seems|it appears)\\b\\s*',
  'gi'
);

const LEADERS = new RegExp(
  '^(?:i\'?ll|i will|i can|i\'?d|you can|we will|we can|let me|let\'?s)\\s+',
  'gim'
);

const ARTICLES = /\b(?:a|an|the)\s+(?=[a-z])/gi;

// Upper bound on sentinel-restore passes. Patterns can nest (e.g. the path
// rule swallows `STARTER/BUSINESS`, then the function-call rule swallows the
// resulting `type ( 0 )`), so a single pass leaves the inner sentinel intact.
// Pass count grows with nesting depth, never with input size — 8 is far past
// any real-world depth and bounds pathological input.
const MAX_RESTORE_PASSES = 8;

// Tokens we won't touch even if they sit inside prose.
const PROTECTED_PATTERNS = [
  /```[\s\S]*?```/g,                          // fenced code
  /`[^`\n]+`/g,                               // inline code
  /\bhttps?:\/\/\S+/gi,                       // URLs
  /\b[\w.-]*[\/\\][\w.\/\\\-]+/g,             // paths with / or \
  /\b[A-Z][A-Za-z0-9]*(?:_[A-Z][A-Za-z0-9]*)+\b/g, // CONST_CASE
  /\b\w+\.\w+(?:\.\w+)*\(\)?/g,               // dotted.method or pkg.fn()
  /[A-Za-z_][A-Za-z0-9_]*\s*\([^)]*\)/g,      // function calls
  /\b\d+\.\d+\.\d+\b/g,                       // version numbers
];

function withProtectedSegments(text, transform) {
  // Replace every protected match with a sentinel, transform the rest, then
  // splice the originals back in.
  const segments = [];
  let working = text;
  for (const re of PROTECTED_PATTERNS) {
    working = working.replace(re, m => {
      const i = segments.length;
      segments.push(m);
      return `${i}`;
    });
  }
  let out = transform(working);
  // Restore iteratively: a single pass leaves inner sentinels unresolved when
  // patterns nested at protect time (#444).
  const sentinel = /(\d+)/g;
  for (let pass = 0; pass < MAX_RESTORE_PASSES; pass++) {
    sentinel.lastIndex = 0;
    if (!sentinel.test(out)) break;
    out = out.replace(/(\d+)/g, (_, i) => segments[+i]);
  }
  return out;
}

function compressProse(text) {
  let s = text;
  s = s.replace(LEADERS, '');
  s = s.replace(PLEASANTRIES, '');
  s = s.replace(HEDGES, '');
  s = s.replace(FILLERS, '');
  s = s.replace(ARTICLES, '');
  // Collapse repeated whitespace introduced by removals.
  s = s.replace(/[ \t]{2,}/g, ' ');
  s = s.replace(/\s+([,.;:!?])/g, '$1');
  s = s.replace(/\n{3,}/g, '\n\n');
  // Capitalize the first letter of each sentence we may have left lowercase.
  s = s.replace(/(^|[.!?]\s+)([a-z])/g, (_, pre, ch) => pre + ch.toUpperCase());
  return s.trim();
}

function compress(text, _opts) {
  if (typeof text !== 'string' || text.length === 0) {
    return { compressed: text, before: 0, after: 0 };
  }
  const before = text.length;
  const compressed = withProtectedSegments(text, compressProse);
  return { compressed, before, after: compressed.length };
}

// Walk a JSON-RPC payload and compress every `description` field in place.
// Used by the proxy on tools/list, prompts/list, resources/list responses.
function compressDescriptionsInPlace(obj, fieldNames) {
  const fields = new Set(fieldNames || ['description']);
  if (!obj || typeof obj !== 'object') return;
  if (Array.isArray(obj)) {
    for (const item of obj) compressDescriptionsInPlace(item, [...fields]);
    return;
  }
  for (const [key, val] of Object.entries(obj)) {
    if (fields.has(key) && typeof val === 'string') {
      obj[key] = compress(val).compressed;
    } else if (val && typeof val === 'object') {
      compressDescriptionsInPlace(val, [...fields]);
    }
  }
}

module.exports = { compress, compressDescriptionsInPlace, withProtectedSegments };


-- --------------------------------------------------------------------------------------------------------------
-- SCHEMA FOR MAESTRO WORKFLOW RELATED DAOs
-- --------------------------------------------------------------------------------------------------------------
CREATE SEQUENCE IF NOT EXISTS workflow_id_seq START 10000 INCREMENT 1;                         -- starting from 10K

CREATE TABLE IF NOT EXISTS maestro_workflow (   -- table to lock and store mutable workflow definition and instance info
  workflow_id           TEXT NOT NULL COLLATE "C",
  internal_id           INT8 DEFAULT nextval('workflow_id_seq') NOT NULL,       -- internal unique sequence id
  active_version_id     INT8 DEFAULT 0 CHECK (active_version_id >= 0),          -- 0 means inactive
  activate_ts           TIMESTAMPTZ,
  activated_by          JSONB,
  properties_snapshot   JSONB CHECK (properties_snapshot IS NOT NULL),          -- current properties snapshot
  latest_version_id     INT8 DEFAULT 1 CHECK (latest_version_id > 0),           -- latest workflow version id
  latest_instance_id    INT8 DEFAULT 0 NOT NULL CHECK (latest_instance_id >= 0),-- latest workflow instance id
  modify_ts             TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP NOT NULL,         -- last modified timestamp
  PRIMARY KEY (workflow_id)
);

CREATE TABLE IF NOT EXISTS maestro_workflow_timeline (  -- table to store the workflow changes
  workflow_id   TEXT NOT NULL COLLATE "C",
  change_event  JSONB NOT NULL,
  hash_id       INT8 NOT NULL,
  create_ts     TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP NOT NULL,
  PRIMARY KEY (workflow_id, create_ts, hash_id)
);

CREATE TABLE IF NOT EXISTS maestro_workflow_version (   -- table of workflow version history, immutable
  workflow_id   TEXT NOT NULL COLLATE "C",
  version_id    INT8 NOT NULL CHECK (version_id > 0),
  metadata      JSONB NOT NULL,
  definition    JSON NOT NULL,
  trigger_uuids JSONB,
  create_ts     TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP NOT NULL, -- processing delay is the time diff between it and metadata create_time
  PRIMARY KEY (workflow_id, version_id)
);

CREATE TABLE IF NOT EXISTS maestro_workflow_properties (    -- table of properties changes and previous snapshot, immutable
  workflow_id           TEXT NOT NULL COLLATE "C",
  create_time           INT8 NOT NULL,      -- is treated as version (not strictly increasing), add version if needed
  author                JSONB NOT NULL,    -- author for the properties_changes
  properties_changes    JSONB NOT NULL,     -- properties changes (delta) for the workflow level settings
  previous_snapshot     JSONB,              -- properties snapshot just before the changes, used for reverting
  PRIMARY KEY (workflow_id, create_time)
);

CREATE TABLE IF NOT EXISTS maestro_workflow_deleted (   -- table to store the deleted workflow basic info for auditing
  workflow_id   TEXT NOT NULL COLLATE "C",       -- workflow id can be re-used
  internal_id   INT8 NOT NULL,     -- make the record unique
  workflow      JSONB NOT NULL,                                         -- a copy of data deleted from maestro_workflow
  create_ts     TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP NOT NULL,         -- the moment the row is inserted
  stage         TEXT DEFAULT 'DELETING_VERSIONS' NOT NULL,            -- stage enum name
  timeline      TEXT[] NOT NULL,                                      -- delete timeline info, e.g. who deletes it, etc.
  modify_ts     TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP NOT NULL,         -- last modified timestamp
  PRIMARY KEY (workflow_id, internal_id)
);


package io.unitycatalog.server.delta.serde;

import com.fasterxml.jackson.databind.BeanDescription;
import com.fasterxml.jackson.databind.DeserializationConfig;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.SerializationConfig;
import com.fasterxml.jackson.databind.deser.BeanDeserializerModifier;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.ser.BeanSerializerModifier;
import io.unitycatalog.server.delta.model.DeltaDataType;
import io.unitycatalog.server.delta.model.DeltaDecimalType;
import io.unitycatalog.server.delta.model.DeltaPrimitiveType;
import io.unitycatalog.server.delta.serde.internal.DeltaTypeDeserializer;
import io.unitycatalog.server.delta.serde.internal.DeltaTypeSerializer;

/**
 * Jackson module that adds Delta's string-or-object wire format for {@link DeltaDataType}.
 *
 *
Register on any {@code ObjectMapper} that (de)serializes UC Delta API models:
 *
 *
{@code
 * mapper.registerModule(new DeltaDataTypeModule());
 * }

 *
 *
Once registered, the mapper reads/writes bare strings (e.g. {@code "long"}, {@code
 * "decimal(10,2)"}) for {@link DeltaPrimitiveType} / {@link DeltaDecimalType}, while still using
 * the default object form for {@code array}, {@code map}, and {@code struct}.
 */
public class DeltaDataTypeModule extends SimpleModule {
  @SuppressWarnings("unchecked")
  public DeltaDataTypeModule() {
    super("DeltaDataTypeModule");
    setDeserializerModifier(
        new BeanDeserializerModifier() {
          @Override
          public JsonDeserializer<?> modifyDeserializer(
              DeserializationConfig config,
              BeanDescription desc,
              JsonDeserializer<?> deserializer) {
            if (desc.getBeanClass() == DeltaDataType.class) {
              return new DeltaTypeDeserializer(deserializer);
            }
            return deserializer;
          }
        });
    setSerializerModifier(
        new BeanSerializerModifier() {
          @Override
          public JsonSerializer<?> modifySerializer(
              SerializationConfig config, BeanDescription desc, JsonSerializer<?> serializer) {
            Class<?> cls = desc.getBeanClass();
            if (cls == DeltaDataType.class
                || cls == DeltaPrimitiveType.class
                || cls == DeltaDecimalType.class) {
              return new DeltaTypeSerializer((JsonSerializer<Object>) serializer);
            }
            return serializer;
          }
        });
  }
}


# pyright: reportMissingImports=false
import asyncio
import os
import sys
from dataclasses import dataclass

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from dotenv import load_dotenv

load_dotenv()

# Third-party imports
import gradio as gr  # type: ignore
from rich.console import Console
from rich.panel import Panel
from rich.text import Text

# Local module imports
from browser_use import Agent, ChatOpenAI

@dataclass
class ActionResult:
	is_done: bool
	extracted_content: str | None
	error: str | None
	include_in_memory: bool

@dataclass
class AgentHistoryList:
	all_results: list[ActionResult]
	all_model_outputs: list[dict]

def parse_agent_history(history_str: str) -> None:
	console = Console()

	# Split the content into sections based on ActionResult entries
	sections = history_str.split('ActionResult(')

	for i, section in enumerate(sections[1:], 1):  # Skip first empty section
		# Extract relevant information
		content = ''
		if 'extracted_content=' in section:
			content = section.split('extracted_content=')[1].split(',')[0].strip("'")

		if content:
			header = Text(f'Step {i}', style='bold blue')
			panel = Panel(content, title=header, border_style='blue')
			console.print(panel)
			console.print()

	return None

async def run_browser_task(
	task: str,
	api_key: str,
	model: str = 'gpt-4.1',
	headless: bool = True,
) -> str:
	if not api_key.strip():
		return 'Please provide an API key'

	os.environ['OPENAI_API_KEY'] = api_key

	try:
		agent = Agent(
			task=task,
			llm=ChatOpenAI(model='gpt-4.1-mini'),
		)
		result = await agent.run()
		#  TODO: The result could be parsed better
		return str(result)
	except Exception as e:
		return f'Error: {str(e)}'

def create_ui():
	with gr.Blocks(title='Browser Use GUI') as interface:
		gr.Markdown('# Browser Use Task Automation')

		with gr.Row():
			with gr.Column():
				api_key = gr.Textbox(label='OpenAI API Key', placeholder='sk-...', type='password')
				task = gr.Textbox(
					label='Task Description',
					placeholder='E.g., Find flights from New York to London for next week',
					lines=3,
				)
				model = gr.Dropdown(choices=['gpt-4.1-mini', 'gpt-5', 'o3', 'gpt-5-mini'], label='Model', value='gpt-4.1-mini')
				headless = gr.Checkbox(label='Run Headless', value=False)
				submit_btn = gr.Button('Run Task')

			with gr.Column():
				output = gr.Textbox(label='Output', lines=10, interactive=False)

		submit_btn.click(
			fn=lambda *args: asyncio.run(run_browser_task(*args)),
			inputs=[task, api_key, model, headless],
			outputs=output,
		)

	return interface

if __name__ == '__main__':
	demo = create_ui()
	demo.launch()


// Unit tests for the argv parser embedded in bin/install.js.
// We don't import parseArgs (it's not exported) — instead we shell out to the
// installer with --help / --list / unknown flags and assert the framing.
// For deeper coverage of flag-resolution semantics, exec --dry-run --list and
// check the rendered defaults.

import { test } from 'node:test';
import assert from 'node:assert/strict';
import { spawnSync } from 'node:child_process';
import path from 'node:path';
import { fileURLToPath } from 'node:url';

const HERE = path.dirname(fileURLToPath(import.meta.url));
const INSTALLER = path.resolve(HERE, '..', '..', 'bin', 'install.js');

function run(...args) {
  return spawnSync('node', [INSTALLER, ...args], { encoding: 'utf8' });
}

test('--help prints usage and exits 0', () => {
  const r = run('--help');
  assert.equal(r.status, 0);
  assert.match(r.stdout, /USAGE/);
  assert.match(r.stdout, /--with-hooks/);
});

test('--list prints provider matrix', () => {
  const r = run('--list');
  assert.equal(r.status, 0);
  assert.match(r.stdout, /caveman provider matrix/);
  assert.match(r.stdout, /claude\b/);
  assert.match(r.stdout, /gemini\b/);
  assert.match(r.stdout, /antigravity\b.*\(soft\)/);
});

test('unknown flag exits 2 with error', () => {
  const r = run('--bogus');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /unknown flag/);
});

test('--all + --minimal mutually exclusive', () => {
  const r = run('--all', '--minimal');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /mutually exclusive/);
});

test('--only without arg fails', () => {
  const r = run('--only');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /--only requires an argument/);
});

test('--config-dir without arg fails', () => {
  const r = run('--config-dir');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /--config-dir requires a path/);
});

test('--config-dir followed by another flag fails', () => {
  const r = run('--config-dir', '--all');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /--config-dir requires a path/);
});

test('aider alias rewrites to aider-desk in dry-run output', () => {
  const r = run('--dry-run', '--only', 'aider', '--non-interactive', '--config-dir', '/tmp/__cm_alias_test');
  // No detection means no install lines, but the script should not crash.
  assert.equal(r.status, 0);
});

test('--only with unknown agent id exits 2', () => {
  const r = run('--only', 'definitely-not-an-agent', '--non-interactive');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /unknown agent: definitely-not-an-agent/);
  assert.match(r.stderr, /caveman --list/);
});

test('--only known id passes argv validation', () => {
  // Dry-run + --only claude exits 0 even if the claude binary isn't on PATH.
  const r = run('--dry-run', '--only', 'claude', '--non-interactive', '--config-dir', '/tmp/__cm_only_test');
  assert.equal(r.status, 0);
});

test('--config-dir expands ~ to home directory', async () => {
  // Pass `~/cm-test-…` and assert the dry-run plan resolves it relative to $HOME.
  // Use a unique suffix so the assertion is unambiguous.
  const suffix = `cm-test-${process.pid}`;
  // --with-hooks: since #392/#393 the hooks plan (which echoes the resolved
  // config-dir) is only emitted when the plugin install fails OR hooks are
  // forced. Force it so the path-expansion assertion below has something to
  // match even when the caveman plugin is already installed.
  const r = run('--dry-run', '--only', 'claude', '--with-hooks', '--non-interactive', '--config-dir', `~/${suffix}`);
  assert.equal(r.status, 0);
  // If the literal `~` had survived, we'd see `~/cm-test-…/hooks` in the plan.
  // The fix expands it in parseArgs, so we expect the absolute home path.
  assert.doesNotMatch(r.stdout, /~\/cm-test-/);
  // The plan only includes the hooks dir if claude is detected. Skip the
  // positive assertion when claude isn't on PATH on the runner.
  if (/Claude Code detected/.test(r.stdout)) {
    const { homedir } = await import('node:os');
    assert.match(r.stdout, new RegExp(homedir().replace(/[.*+?^${}()|[\]\\]/g, '\\$&') + '/' + suffix));
  }
});

test('bare -- (POSIX end-of-options) is accepted and ignored', () => {
  // Regression: npx forwarded `--` from `curl|bash -- --only openclaw` to the
  // package, and parseArgs rejected it as an unknown flag. Now we accept it.
  const r = run('--', '--only', 'claude', '--non-interactive', '--dry-run', '--config-dir', '/tmp/__cm_dashdash');
  assert.equal(r.status, 0);
});

test('bare --with-mcp-shrink (no upstream) exits 2 with hint', () => {
  // Regression for issue where --with-mcp-shrink registered a stub MCP entry
  // that crashed on every Claude Code startup. caveman-shrink is a proxy and
  // requires an upstream command — we now refuse the bare flag (#474).
  const r = run('--with-mcp-shrink', '--non-interactive', '--dry-run');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /requires an upstream command/);
  assert.match(r.stderr, /server-filesystem/);
});

test('--with-mcp-shrink followed by another flag (no value) exits 2', () => {
  // The next-token form must distinguish "no value" from "value happens to
  // start with --". A user typing `--with-mcp-shrink --dry-run` clearly
  // forgot the upstream; refuse.
  const r = run('--with-mcp-shrink', '--dry-run', '--non-interactive');
  assert.equal(r.status, 2);
  assert.match(r.stderr, /requires an upstream command/);
});

test('--with-mcp-shrink="<cmd>" registers wrapping that upstream', () => {
  const r = run(
    '--with-mcp-shrink=npx @modelcontextprotocol/server-filesystem /tmp',
    '--only', 'claude', '--dry-run', '--non-interactive',
    '--config-dir', '/tmp/__cm_shrink_test'
  );
  assert.equal(r.status, 0);
  // Dry-run only emits the planned `claude mcp add` line when claude is on
  // PATH (installMcpShrink probes `claude mcp --help` first). Assert the
  // wrapping content only when that line is actually present.
  if (/would run: claude mcp add caveman-shrink/.test(r.stdout)) {
    assert.match(r.stdout, /claude mcp add caveman-shrink .* npx -y caveman-shrink npx @modelcontextprotocol\/server-filesystem \/tmp/);
  }
});

test('--with-mcp-shrink "<cmd>" (space-separated) also accepted', () => {
  const r = run(
    '--with-mcp-shrink', 'npx @modelcontextprotocol/server-filesystem /tmp',
    '--only', 'claude', '--dry-run', '--non-interactive',
    '--config-dir', '/tmp/__cm_shrink_space'
  );
  assert.equal(r.status, 0);
  if (/would run: claude mcp add caveman-shrink/.test(r.stdout)) {
    assert.match(r.stdout, /caveman-shrink npx @modelcontextprotocol\/server-filesystem \/tmp/);
  }
});

test('--all does NOT auto-enable mcp-shrink (no sensible default upstream)', () => {
  const r = run('--all', '--only', 'claude', '--dry-run', '--non-interactive', '--config-dir', '/tmp/__cm_all_no_shrink');
  assert.equal(r.status, 0);
  // Whether or not claude is on PATH, the wiring banner should not appear
  // because withMcpShrink stays false under --all alone.
  assert.doesNotMatch(r.stdout, /wiring caveman-shrink MCP proxy/);
});

test('--help discloses --config-dir scope', () => {
  const r = run('--help');
  assert.equal(r.status, 0);
  // Disclosure: --config-dir does NOT scope third-party CLI invocations.
  // Help text wraps mid-phrase, so collapse whitespace before matching.
  const collapsed = r.stdout.replace(/\s+/g, ' ');
  assert.match(collapsed, /Does NOT scope/);
  assert.match(collapsed, /XDG_CONFIG_HOME/);
  assert.match(collapsed, /OPENCLAW_WORKSPACE/);
});


/*
===============================================================================
Quality Checks
===============================================================================
Script Purpose:
    This script performs various quality checks for data consistency, accuracy,
    and standardization across the 'silver' layer. It includes checks for:
    - Null or duplicate primary keys.
    - Unwanted spaces in string fields.
    - Data standardization and consistency.
    - Invalid date ranges and orders.
    - Data consistency between related fields.

Usage Notes:
    - Run these checks after data loading Silver Layer.
    - Investigate and resolve any discrepancies found during the checks.
===============================================================================
*/

-- ====================================================================
-- Checking 'silver.crm_cust_info'
-- ====================================================================
-- Check for NULLs or Duplicates in Primary Key
-- Expectation: No Results
SELECT
    cst_id,
    COUNT(*)
FROM silver.crm_cust_info
GROUP BY cst_id
HAVING COUNT(*) > 1 OR cst_id IS NULL;

-- Check for Unwanted Spaces
-- Expectation: No Results
SELECT
    cst_key
FROM silver.crm_cust_info
WHERE cst_key != TRIM(cst_key);

-- Data Standardization & Consistency
SELECT DISTINCT
    cst_marital_status
FROM silver.crm_cust_info;

-- ====================================================================
-- Checking 'silver.crm_prd_info'
-- ====================================================================
-- Check for NULLs or Duplicates in Primary Key
-- Expectation: No Results
SELECT
    prd_id,
    COUNT(*)
FROM silver.crm_prd_info
GROUP BY prd_id
HAVING COUNT(*) > 1 OR prd_id IS NULL;

-- Check for Unwanted Spaces
-- Expectation: No Results
SELECT
    prd_nm
FROM silver.crm_prd_info
WHERE prd_nm != TRIM(prd_nm);

-- Check for NULLs or Negative Values in Cost
-- Expectation: No Results
SELECT
    prd_cost
FROM silver.crm_prd_info
WHERE prd_cost < 0 OR prd_cost IS NULL;

-- Data Standardization & Consistency
SELECT DISTINCT
    prd_line
FROM silver.crm_prd_info;

-- Check for Invalid Date Orders (Start Date > End Date)
-- Expectation: No Results
SELECT
    *
FROM silver.crm_prd_info
WHERE prd_end_dt < prd_start_dt;

-- ====================================================================
-- Checking 'silver.crm_sales_details'
-- ====================================================================
-- Check for Invalid Dates
-- Expectation: No Invalid Dates
SELECT
    NULLIF(sls_due_dt, 0) AS sls_due_dt
FROM bronze.crm_sales_details
WHERE sls_due_dt <= 0
    OR LEN(sls_due_dt) != 8
    OR sls_due_dt > 20500101
    OR sls_due_dt < 19000101;

-- Check for Invalid Date Orders (Order Date > Shipping/Due Dates)
-- Expectation: No Results
SELECT
    *
FROM silver.crm_sales_details
WHERE sls_order_dt > sls_ship_dt
   OR sls_order_dt > sls_due_dt;

-- Check Data Consistency: Sales = Quantity * Price
-- Expectation: No Results
SELECT DISTINCT
    sls_sales,
    sls_quantity,
    sls_price
FROM silver.crm_sales_details
WHERE sls_sales != sls_quantity * sls_price
   OR sls_sales IS NULL
   OR sls_quantity IS NULL
   OR sls_price IS NULL
   OR sls_sales <= 0
   OR sls_quantity <= 0
   OR sls_price <= 0
ORDER BY sls_sales, sls_quantity, sls_price;

-- ====================================================================
-- Checking 'silver.erp_cust_az12'
-- ====================================================================
-- Identify Out-of-Range Dates
-- Expectation: Birthdates between 1924-01-01 and Today
SELECT DISTINCT
    bdate
FROM silver.erp_cust_az12
WHERE bdate < '1924-01-01'
   OR bdate > GETDATE();

-- Data Standardization & Consistency
SELECT DISTINCT
    gen
FROM silver.erp_cust_az12;

-- ====================================================================
-- Checking 'silver.erp_loc_a101'
-- ====================================================================
-- Data Standardization & Consistency
SELECT DISTINCT
    cntry
FROM silver.erp_loc_a101
ORDER BY cntry;

-- ====================================================================
-- Checking 'silver.erp_px_cat_g1v2'
-- ====================================================================
-- Check for Unwanted Spaces
-- Expectation: No Results
SELECT
    *
FROM silver.erp_px_cat_g1v2
WHERE cat != TRIM(cat)
   OR subcat != TRIM(subcat)
   OR maintenance != TRIM(maintenance);

-- Data Standardization & Consistency
SELECT DISTINCT
    maintenance
FROM silver.erp_px_cat_g1v2;