ai-teacher/backend/src/main/java/com/aiteacher/topic/TopicSummaryService.java

package com.aiteacher.topic;

import com.aiteacher.book.Book;
import com.aiteacher.book.BookRepository;
import com.aiteacher.book.BookStatus;
import com.aiteacher.book.NoKnowledgeSourceException;
import com.aiteacher.document.FigureEntity;
import com.aiteacher.document.SectionEntity;
import com.aiteacher.retrieval.NeurosurgeryRetriever;
import com.aiteacher.retrieval.RetrievalResult;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.stereotype.Service;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.UUID;

@Service
public class TopicSummaryService {

    private static final Logger log = LoggerFactory.getLogger(TopicSummaryService.class);

    private static final String SYSTEM_PROMPT = """
        You are an expert neurosurgery educator. Your role is to provide accurate,
        clinically relevant summaries based ONLY on the content retrieved from the
        uploaded medical textbooks. Do not use any knowledge outside the provided context.

        When answering:
        - Structure your response clearly with key points
        - Cite claims using ONLY the reference labels provided in the context (e.g. [S1], [F2]).
          Do not invent page numbers, section titles, or labels not present in the CONTEXT block.
        - If the retrieved context does not contain sufficient information on the topic,
          explicitly state: "The uploaded books do not contain sufficient information on this topic."
        - Never hallucinate or fabricate clinical information
        """;

    private final ChatClient chatClient;
    private final BookRepository bookRepository;
    private final NeurosurgeryRetriever retriever;
    private final TopicSummaryRepository summaryRepository;
    private final ObjectMapper objectMapper;

    public TopicSummaryService(ChatClient chatClient,
                                BookRepository bookRepository,
                                NeurosurgeryRetriever retriever,
                                TopicSummaryRepository summaryRepository,
                                ObjectMapper objectMapper) {
        this.chatClient = chatClient;
        this.bookRepository = bookRepository;
        this.retriever = retriever;
        this.summaryRepository = summaryRepository;
        this.objectMapper = objectMapper;
    }

    public TopicSummaryResponse generateSummary(Topic topic) {
        List<Book> readyBooks = bookRepository.findAll().stream()
            .filter(b -> b.getStatus() == BookStatus.READY)
            .toList();

        if (readyBooks.isEmpty()) {
            throw new NoKnowledgeSourceException(
                "No books are available as knowledge sources. Please upload and process at least one book.");
        }

        String question = buildQuestion(topic);

        List<SectionEntity> allSections = new ArrayList<>();
        List<FigureEntity> allFigures = new ArrayList<>();
        for (Book book : readyBooks) {
            RetrievalResult result = retriever.retrieve(question, book.getId());
            allSections.addAll(result.parentSections());
            allFigures.addAll(result.figures());
        }

        log.debug("Topic summary for '{}': {} sections, {} figures retrieved",
            topic.getName(), allSections.size(), allFigures.size());

        String contextPrompt = buildContextPrompt(question, allSections, allFigures);
        String summary = chatClient.prompt()
            .system(SYSTEM_PROMPT)
            .user(contextPrompt)
            .call()
            .content();

        List<TopicSummaryResponse.SourceReference> sources = buildSources(allSections, allFigures, readyBooks);
        Instant generatedAt = Instant.now();

        int summaryNumber = (int) summaryRepository.countByTopicId(topic.getId()) + 1;
        String sourcesJson = serializeSources(sources);
        TopicSummaryEntity entity = new TopicSummaryEntity(
            topic.getId(), summaryNumber, summary, sourcesJson, generatedAt);
        entity = summaryRepository.save(entity);

        return new TopicSummaryResponse(
            entity.getId(),
            summaryNumber,
            topic.getId(),
            topic.getName(),
            summary,
            sources,
            generatedAt
        );
    }

    public List<SavedSummaryItem> listSummaries(String topicId) {
        return summaryRepository.findByTopicIdOrderBySummaryNumberAsc(topicId).stream()
            .map(e -> new SavedSummaryItem(e.getId(), e.getSummaryNumber(), e.getGeneratedAt()))
            .toList();
    }

    public TopicSummaryResponse getSummary(UUID summaryId) {
        TopicSummaryEntity entity = summaryRepository.findById(summaryId)
            .orElseThrow(() -> new NoSuchElementException("Summary not found."));

        List<TopicSummaryResponse.SourceReference> sources = deserializeSources(entity.getSourcesJson());

        return new TopicSummaryResponse(
            entity.getId(),
            entity.getSummaryNumber(),
            entity.getTopicId(),
            entity.getTopicId(),
            entity.getSummary(),
            sources,
            entity.getGeneratedAt()
        );
    }

    private String buildQuestion(Topic topic) {
        return String.format(
            "Provide a comprehensive educational summary of the following neurosurgery topic: " +
            "%s. Topic description: %s. " +
            "Include key concepts, clinical considerations, and important details that a neurosurgeon should know.",
            topic.getName(), topic.getDescription()
        );
    }

    private String buildContextPrompt(String question,
                                      List<SectionEntity> sections,
                                      List<FigureEntity> figures) {
        StringBuilder sb = new StringBuilder();

        if (!sections.isEmpty()) {
            sb.append("CONTEXT:\n\n");
            for (int i = 0; i < sections.size(); i++) {
                SectionEntity s = sections.get(i);
                sb.append("[S").append(i + 1).append("] ")
                  .append(s.getTitle()).append(", p.").append(s.getPageStart()).append("\n");
                sb.append(s.getFullText()).append("\n\n");
            }
        }

        if (!figures.isEmpty()) {
            sb.append("AVAILABLE FIGURES:\n");
            for (int i = 0; i < figures.size(); i++) {
                FigureEntity f = figures.get(i);
                sb.append("[F").append(i + 1).append("] ")
                  .append(f.getLabel() != null ? f.getLabel() : "Figure")
                  .append(" (p.").append(f.getPage()).append("): ")
                  .append(f.getCaption() != null ? f.getCaption() : "")
                  .append("\n");
            }
            sb.append("\n");
        }

        sb.append("QUESTION:\n").append(question);
        return sb.toString();
    }

    private List<TopicSummaryResponse.SourceReference> buildSources(List<SectionEntity> sections,
                                                                      List<FigureEntity> figures,
                                                                      List<Book> readyBooks) {
        List<TopicSummaryResponse.SourceReference> sources = new ArrayList<>();

        for (SectionEntity s : sections) {
            Book book = readyBooks.stream()
                .filter(b -> b.getId().equals(s.getBookId()))
                .findFirst()
                .orElse(null);
            String title = book != null ? book.getTitle() : "Book";
            String bookId = book != null ? book.getId().toString() : null;
            sources.add(new TopicSummaryResponse.SourceReference(bookId, title, s.getPageStart()));
        }

        for (FigureEntity f : figures) {
            Book book = readyBooks.stream()
                .filter(b -> b.getId().equals(f.getBookId()))
                .findFirst()
                .orElse(null);
            String title = book != null ? book.getTitle() : "Book";
            String bookId = book != null ? book.getId().toString() : null;
            sources.add(new TopicSummaryResponse.SourceReference(bookId, title, f.getPage()));
        }

        return sources.stream().distinct().toList();
    }

    private String serializeSources(List<TopicSummaryResponse.SourceReference> sources) {
        try {
            return objectMapper.writeValueAsString(sources);
        } catch (JsonProcessingException e) {
            log.warn("Failed to serialize sources, storing empty array", e);
            return "[]";
        }
    }

    private List<TopicSummaryResponse.SourceReference> deserializeSources(String json) {
        try {
            return objectMapper.readValue(json,
                objectMapper.getTypeFactory().constructCollectionType(
                    List.class, TopicSummaryResponse.SourceReference.class));
        } catch (JsonProcessingException e) {
            log.warn("Failed to deserialize sources from stored JSON", e);
            return List.of();
        }
    }
}