package ankurm.ai21;

import java.util.List;
import java.util.Map;

import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.AssistantMessage.ToolCall;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.part.MessagePart;
import org.springframework.ai.chat.messages.part.OpaquePayload;
import org.springframework.ai.chat.messages.part.ReasoningPart;
import org.springframework.ai.chat.messages.part.StreamingParts;
import org.springframework.ai.chat.messages.part.TextPart;
import org.springframework.ai.chat.messages.part.ToolCallPart;

/** Spring AI 2.1.0-M1: build an assistant turn as ordered parts, then inspect it. */
public class PartsDemo {

	public static void main(String[] args) {
		// 1. One assistant turn: reasoning -> tool call -> text, in the order the provider produced it.
		OpaquePayload signature = new OpaquePayload("openai", "reasoning.encrypted_content", "gAAAA-not-a-real-signature");
		AssistantMessage turn = AssistantMessage.builder()
			.part(new ReasoningPart("Need the weather first.", "Plans a lookup", signature, Map.of()))
			.part(ToolCallPart.of(new ToolCall("call_1", "function", "getWeather", "{\"city\":\"Mumbai\"}")))
			.part(TextPart.of("Checking the weather for Mumbai."))
			.build();

		System.out.println("== 1. parts in provider order");
		for (MessagePart p : turn.getParts()) {
			System.out.println(p.getClass().getSimpleName());
		}
		System.out.println("getText()      = " + turn.getText());
		System.out.println("getToolCalls() = " + turn.getToolCalls().size());
		System.out.println("getReasoning() = " + turn.getReasoning().size());

		// 2. Streaming: a partial part carries its index; complete() marks it finished.
		System.out.println("== 2. streaming markers");
		MessagePart d1 = StreamingParts.partial(TextPart.of("Hel"), 0);
		MessagePart d2 = StreamingParts.partial(TextPart.of("lo"), 0);
		MessagePart done = StreamingParts.complete(TextPart.of("Hello"), 0);
		for (MessagePart p : List.of(d1, d2, done)) {
			System.out.println("partIndex=" + StreamingParts.partIndex(p) + " partial=" + StreamingParts.isPartial(p)
					+ " attributes=" + p.attributes());
		}
		System.out.println("strip(done).attributes = " + StreamingParts.strip(done).attributes());

		// 3. Replaying a provider signature: only to the provider that issued it.
		ReasoningPart signed = turn.getReasoning().get(0);
		System.out.println("== 3. replayableTo");
		System.out.println("signed   -> openai    : " + signed.replayableTo("openai"));
		System.out.println("signed   -> anthropic : " + signed.replayableTo("anthropic"));
		System.out.println("unsigned -> openai    : " + ReasoningPart.of("plain text reasoning").replayableTo("openai"));

		// 4. Does the in-memory chat memory repository keep the reasoning part across turns?
		InMemoryChatMemoryRepository repo = new InMemoryChatMemoryRepository();
		repo.saveAll("c1", List.of(new UserMessage("Weather in Mumbai?"), turn));
		List<Message> loaded = repo.findByConversationId("c1");
		AssistantMessage back = (AssistantMessage) loaded.get(1);
		System.out.println("== 4. in-memory repository round trip");
		System.out.println("reasoning parts after reload = " + back.getReasoning().size());
		System.out.println("payload kept                 = " + back.getReasoning().get(0).payload());
	}
}