package com.acumenvelocity.ath.mt.confidence;
   
   import com.acumenvelocity.ath.common.AthUtil;
   import com.acumenvelocity.ath.common.Const;
   import com.acumenvelocity.ath.common.Log;
   import com.google.gson.*;
   
   import java.io.*;
   import java.nio.charset.StandardCharsets;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.util.*;
  import java.util.concurrent.*;
  import java.util.stream.Collectors;
  
  public class VertexQualityEstimator implements AutoCloseable {
  
    private final ExecutorService executor = Executors.newFixedThreadPool(
        Math.max(2, Runtime.getRuntime().availableProcessors() - 1));
  
    private final Gson gson = new Gson();
  
    private static final String PYTHON_CMD = System.getProperty("python.cmd", "python3");
    private static final String SCRIPT_RESOURCE_PATH = "/eval_qe_production.py";
  
    // Temp script extracted from classpath — created once
    private final Path extractedScriptPath;
  
    public VertexQualityEstimator() throws IOException {
      this.extractedScriptPath = extractScriptFromClasspath();
    }
  
    private Path extractScriptFromClasspath() throws IOException {
      InputStream is = getClass().getResourceAsStream(SCRIPT_RESOURCE_PATH);
  
      if (is == null) {
        throw new IOException("Python script not found in classpath: " + SCRIPT_RESOURCE_PATH +
            " — make sure eval_qe_production.py is in src/main/resources/");
      }
  
      Path tempFile = Files.createTempFile("eval_qe_", ".py");
      tempFile.toFile().deleteOnExit();
  
      try (OutputStream os = Files.newOutputStream(tempFile)) {
        is.transferTo(os);
      }
  
      if (!tempFile.toFile().setExecutable(true)) {
        Log.warn(getClass(), "Could not set executable flag on temp Python script: {}", tempFile);
      }
  
      Log.info(getClass(), "Successfully extracted Python QE script to: {}", tempFile);
      return tempFile;
    }
  
    public VertexDataStructures.TranslationComparisonResult evaluateMultiModelTranslations(
        List<CommonDataStructures.SegmentTranslations> segments,
        String sourceLang, String targetLang) throws IOException {
  
      List<CompletableFuture<VertexDataStructures.ScoredSegmentTranslations>> futures = new ArrayList<>();
  
      for (CommonDataStructures.SegmentTranslations seg : segments) {
        futures.add(CompletableFuture.supplyAsync(() -> evaluateSegment(seg), executor));
      }
  
      List<VertexDataStructures.ScoredSegmentTranslations> scored = futures.stream()
          .map(CompletableFuture::join)
          .filter(Objects::nonNull)
          .collect(Collectors.toList());
  
      return new VertexDataStructures.TranslationComparisonResult(scored);
    }
  
    public VertexDataStructures.ScoredSegmentTranslations evaluateSegment(
        CommonDataStructures.SegmentTranslations seg) {
      try {
        JsonArray payload = new JsonArray();
        List<String> modelIds = new ArrayList<>();
  
        for (Map.Entry<String, String> e : seg.getTranslations().entrySet()) {
          JsonObject obj = new JsonObject();
          obj.addProperty("prompt", seg.getSourceText());
          obj.addProperty("response", e.getValue());
          payload.add(obj);
          modelIds.add(e.getKey());
        }
  
        ProcessBuilder pb = new ProcessBuilder(PYTHON_CMD, extractedScriptPath.toString());
        pb.environment().put("GOOGLE_APPLICATION_CREDENTIALS", Const.ATH_GCP_SECRET_FILE);
  
        Process p = pb.start();
  
        try (OutputStream os = p.getOutputStream()) {
          os.write(payload.toString().getBytes(StandardCharsets.UTF_8));
          os.flush();
        }
  
        // Read all output
        StringBuilder fullOutput = new StringBuilder();
 
       try (BufferedReader reader = new BufferedReader(
           new InputStreamReader(p.getInputStream(), StandardCharsets.UTF_8))) {
         String line;
 
         while ((line = reader.readLine()) != null) {
           fullOutput.append(line).append("\n");
         }
       }
 
       StringBuilder stderr = new StringBuilder();
 
       try (BufferedReader errorReader = new BufferedReader(
           new InputStreamReader(p.getErrorStream(), StandardCharsets.UTF_8))) {
         String line;
 
         while ((line = errorReader.readLine()) != null) {
           stderr.append(line).append("\n");
         }
       }
 
       int exitCode = p.waitFor();
 
       if (exitCode != 0) {
         Log.error(getClass(),
             "Python QE failed for segment {}, exitCode={}\nSTDERR:\n{}\nSTDOUT:\n{}",
             seg.getSegmentId(), exitCode,
             stderr.length() > 0 ? stderr.toString().trim() : "(empty)",
             fullOutput.length() > 0 ? fullOutput.toString().trim() : "(empty)");
 
         return null;
       }
 
       // EXTRACT ONLY THE LINE THAT LOOKS LIKE JSON
       String jsonLine = fullOutput.toString().lines()
           .filter(line -> line.trim().startsWith("{") && line.trim().endsWith("}"))
           .findFirst()
           .orElse(null);
 
       if (jsonLine == null) {
         Log.error(getClass(), "No valid JSON found in Python output for segment {}\nOUTPUT:\n{}",
             seg.getSegmentId(), fullOutput.toString().trim());
 
         return null;
       }
 
       JsonObject result = gson.fromJson(jsonLine, JsonObject.class);
       List<VertexDataStructures.ScoredTranslation> list = new ArrayList<>();
 
       for (int i = 0; i < modelIds.size(); i++) {
         String key = "alt_" + i;
         JsonElement altElem = result.get(key);
 
         if (altElem == null || !altElem.isJsonObject()) {
           Log.warn(getClass(), "Missing alt_{} in result for segment {}", i, seg.getSegmentId());
           continue;
         }
 
         JsonObject altObj = altElem.getAsJsonObject();
         JsonElement metricxElem = altObj.get("metricx");
 
         if (metricxElem == null || !metricxElem.isJsonPrimitive()) {
           Log.warn(getClass(), "Invalid metricx for alt_{} in segment {}", i, seg.getSegmentId());
           continue;
         }
 
         double metricx = metricxElem.getAsDouble();
         String src = seg.getSourceText();
         String trg = seg.getTranslations().get(modelIds.get(i));
 
         Log.info(getClass(), "{} trl by '{}': '{}' from src: '{}'",
             altObj,
             AthUtil.extractLastSection(modelIds.get(i)),
             trg,
             src);
 
         list.add(new VertexDataStructures.ScoredTranslation(
             modelIds.get(i),
             seg.getTranslations().get(modelIds.get(i)),
             metricx,
             altObj.has("fluency") ? altObj.get("fluency").getAsDouble() : 0.0,
             altObj.has("coherence") ? altObj.get("coherence").getAsDouble() : 0.0,
             altObj.has("groundedness") ? altObj.get("groundedness").getAsDouble() : 0.0,
             HeuristicQualityEstimator.calculateLengthRatio(src, trg),
             HeuristicQualityEstimator.calculateTextSimilarity(src, trg)));
       }
 
       if (list.isEmpty()) {
         Log.warn(getClass(), "No valid scores parsed for segment {}", seg.getSegmentId());
         return null;
       }
 
       return new VertexDataStructures.ScoredSegmentTranslations(
           seg.getSegmentId(), seg.getSourceText(), seg.getSourceTf(), list);
 
     } catch (Exception e) {
       Log.error(getClass(), e, "Unexpected error in Gemini QE for segment {}", seg.getSegmentId());
       return null;
     }
   }
 
   @Override
   public void close() {
     executor.shutdownNow();
 
     if (extractedScriptPath != null && Files.exists(extractedScriptPath)) {
       try {
         Files.deleteIfExists(extractedScriptPath);
 
       } catch (Exception ignored) {
       }
     }
   }
 }