package ankurm.ai21;

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

import javax.sql.DataSource;

import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.part.OpaquePayload;
import org.springframework.ai.chat.messages.part.ReasoningPart;
import org.springframework.ai.chat.messages.part.TextPart;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.DriverManagerDataSource;

/** Same turn as PartsDemo, but through the JDBC repository on in-memory H2. */
public class JdbcMemoryDemo {

	public static void main(String[] args) throws Exception {
		DataSource ds = new DriverManagerDataSource("jdbc:h2:mem:cm;DB_CLOSE_DELAY=-1", "sa", "");
		JdbcTemplate jdbc = new JdbcTemplate(ds);
		org.springframework.jdbc.datasource.init.ScriptUtils.executeSqlScript(ds.getConnection(),
				new org.springframework.core.io.ClassPathResource(
						"org/springframework/ai/chat/memory/repository/jdbc/schema-h2.sql"));
		JdbcChatMemoryRepository repo = JdbcChatMemoryRepository.builder().jdbcTemplate(jdbc).build();

		AssistantMessage turn = AssistantMessage.builder()
			.part(new ReasoningPart("Need the weather first.", "Plans a lookup",
					new OpaquePayload("openai", "reasoning.encrypted_content", "gAAAA-not-a-real-signature"), Map.of()))
			.part(TextPart.of("Checking the weather for Mumbai."))
			.build();
		System.out.println("before save: reasoning parts = " + turn.getReasoning().size());

		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("after reload: reasoning parts = " + back.getReasoning().size());
		System.out.println("after reload: text = " + back.getText());
		System.out.println("stored row    : " + jdbc.queryForList("SELECT type, content FROM SPRING_AI_CHAT_MEMORY"));
	}
}