package com.agent.management.rag.document;
import com.agent.management.rag.api.RagRetriever; import com.agent.management.rag.model.*; import com.agent.management.repository.*; import com.agent.management.service.impl.EmbeddingBridgeClient; import lombok.RequiredArgsConstructor; import org.springframework.stereotype.Service; import java.util.*; import java.util.function.Function; import java.util.stream.Collectors;
@Service @RequiredArgsConstructor
public class DocumentRagRetriever implements RagRetriever {
 private final EmbeddingBridgeClient bridge; private final KbChunkRepository chunks; private final KbDocumentRepository documents;
 public RagSourceType sourceType(){return RagSourceType.DOCUMENT;}
 public RagRetrievalResult retrieve(RagQuery query){
  RagRetrievalResult out=new RagRetrievalResult(); out.setSourceType(sourceType());
  try{
   int topK=query.getTopK()!=null&&query.getTopK()>0?Math.min(query.getTopK(),100):5; String mode=filter(query,"mode","hybrid");
   String expr=documentFilter(query.getSourceIds()); Map<String,Object> response=bridge.retrieve(query.getQuery(),topK,mode,expr);
   List<Map<String,Object>> hits=maps(response.get("hits"));
   List<String> ids=hits.stream().map(h->text(h.get("chunk_id"))).filter(s->!s.isBlank()).toList();
   Map<String,com.agent.management.model.entity.KbChunk> found=chunks.findByChunkIdIn(ids).stream().collect(Collectors.toMap(com.agent.management.model.entity.KbChunk::getChunkId,Function.identity(),(a,b)->a));
   for(Map<String,Object> hit:hits) out.getEvidences().add(evidence(hit,found.get(text(hit.get("chunk_id"))),mode));
   out.getDiagnostics().put("mode",mode); out.getDiagnostics().put("topK",topK);
  }catch(Exception e){out.getDiagnostics().put("error",e.getMessage());}
  return out;
 }
 private RagEvidence evidence(Map<String,Object> hit, com.agent.management.model.entity.KbChunk chunk,String mode){
  String chunkId=text(hit.get("chunk_id")); String title=text(hit.get("filename")); String sourceId=text(hit.get("document_id"));
  if(chunk!=null){sourceId=String.valueOf(chunk.getDocumentId()); title=documents.findById(chunk.getDocumentId()).map(com.agent.management.model.entity.KbDocument::getName).orElse(title);}
  RagEvidence e=new RagEvidence(); e.setId(text(hit.get("id"))); e.setSourceType(sourceType()); e.setEvidenceType("DOCUMENT_CHUNK"); e.setTitle(title); e.setContent(text(hit.get("text"))); e.setScore(number(hit.get("score"))); e.setSourceName(title);
  RagSourceRef ref=new RagSourceRef(); ref.setSourceId(sourceId); ref.setSourceName(title); ref.setLocator(Map.of("chunkId",chunkId,"pageNumber",value(hit.get("page_number"),0))); e.setSourceRef(ref);
  Map<String,Object> payload=new LinkedHashMap<>(); payload.put("chunkId",chunkId); payload.put("parentChunkId",text(hit.get("parent_chunk_id"))); payload.put("rootChunkId",text(hit.get("root_chunk_id"))); payload.put("chunkLevel",value(hit.get("chunk_level"),0)); if("dense".equals(mode))payload.put("denseScore",e.getScore()); else payload.put("hybridScore",e.getScore()); e.setPayload(payload); e.setMetadata(new LinkedHashMap<>(hit)); return e;
 }
 private static String documentFilter(List<String> ids){if(ids==null||ids.isEmpty())return ""; List<Long> safe=new ArrayList<>(); for(String id:ids)try{safe.add(Long.parseLong(id));}catch(Exception ignored){} return safe.isEmpty()?"":"document_id in ["+safe.stream().map(String::valueOf).collect(Collectors.joining(","))+"]";}
 @SuppressWarnings("unchecked") private static List<Map<String,Object>> maps(Object o){return o instanceof List<?> l?l.stream().filter(Map.class::isInstance).map(x->(Map<String,Object>)x).toList():List.of();}
 private static String filter(RagQuery q,String key,String fallback){Object v=q.getFilters()==null?null:q.getFilters().get(key);return v==null?fallback:String.valueOf(v);}
 private static Object value(Object v,Object d){return v==null?d:v;} private static String text(Object v){return v==null?"":String.valueOf(v);} private static double number(Object v){return v instanceof Number n?n.doubleValue():0;}
}
