package com.acumenvelocity.ath.openai;
   
   import java.util.ArrayList;
   import java.util.Collections;
   import java.util.List;
   
   import com.acumenvelocity.ath.common.AlignmentData.CombinedAlignmentInput;
   import com.acumenvelocity.ath.common.AlignmentData.CombinedAlignmentOutput;
   import com.acumenvelocity.ath.common.Const;
  import com.acumenvelocity.ath.common.JacksonUtil;
  import com.acumenvelocity.ath.common.Log;
  import com.acumenvelocity.ath.model.InlineCodeRef;
  import com.acumenvelocity.ath.model.LayeredSegment;
  import com.fasterxml.jackson.core.type.TypeReference;
  import com.openai.client.OpenAIClient;
  import com.openai.client.okhttp.OpenAIOkHttpClient;
  import com.openai.models.chat.completions.ChatCompletion;
  import com.openai.models.chat.completions.ChatCompletionCreateParams;
  
  import net.sf.okapi.common.Util;
  
  public class OpenAi {
  
    private static final String CODES_REINSERTION_PROMPT = "For every object in the input array, "
        + "align words and phrases in its trgText against its srcText. "
        + "Reinsert source inline codes into trgCodes positioning them around aligned target words. "
        + "CRITICAL REQUIREMENTS: "
        + "1. WORD BOUNDARIES: Codes MUST be placed only at word boundaries (before the first "
        + "character or after the last character of a word), NEVER in the middle of a word. "
        + "A code splitting a word (e.g., 'о<code>т</code>ображать') is a CRITICAL ERROR. "
        + "2. PAIRED CODES: Two codes with the same id are paired (OPENING and CLOSING tagType). "
        + "If paired codes wrap only whitespace in srcText, they must wrap only whitespace in "
        + "trgText. If they wrap text in srcText, they must wrap text in trgText. Wrapping text "
        + "in target when source wraps only spaces is a SEVERE ERROR. "
        + "3. NESTING: Always preserve well-formed nesting of paired codes with no overlapping. "
        + "4. UNPAIRED CODES: If a code has no pair in the segment, treat it as an isolated "
        + "placeholder (its pair is in another segment) and place it at the nearest appropriate "
        + "word boundary. "
        + "5. ORDER: Preserve the codes order wherever possible. "
        + "VALIDATION: Before finalizing positions, verify each code is at a word boundary "
        + "(adjacent to space, punctuation, or segment start/end), not inside a word. "
        + "Input: JSON array of segments. "
        + "Output: array of trgCodes JSON arrays, with length equal to the input array length. "
        + "Return ONLY the JSON array, no explanations, no markdown, no preamble. This is an "
        + "API call with strict response schema - make all decisions autonomously.";
  
    private static final String COMBINED_ALIGNMENT_PROMPT = "You are a bilingual document alignment "
        + "expert. Perform TWO-STAGE alignment:\n\n"
        + "STAGE 1 - PARAGRAPH ALIGNMENT:\n"
        + "- Align sourceParagraphs with targetParagraphs\n"
        + "- Documents may have different paragraph structures (split, merged, reordered)\n"
        + "- Use context, segment previews, and content similarity\n"
        + "- Types: MATCH (1:1), MULTI_MATCH (many:many), DELETED, INSERTED\n\n"
  
        + "STAGE 2 - SENTENCE ALIGNMENT (within each paragraph pair):\n"
        + "- For EACH paragraph alignment, align the sentences within\n"
        + "- Source positions refer to segment positions within SOURCE paragraphs of that pair\n"
        + "- Target positions refer to segment positions within TARGET paragraphs of that pair\n"
        + "- For MULTI_MATCH paragraphs, treat all segments from all source/target paragraphs as continuous lists\n"
        + "- Types: MATCH (1:1), MULTI_MATCH (many:many), DELETED, INSERTED\n\n"
  
        + "CRITICAL REQUIREMENTS:\n"
        + "1. Every paragraph must appear in exactly one paragraph alignment\n"
        + "2. Every segment must appear in exactly one sentence alignment\n"
        + "3. For DELETED: source segment has no target (we'll insert empty target)\n"
        + "4. For INSERTED: target segment has no source (we'll insert empty source)\n"
        + "5. After alignment, source and target will have EQUAL segment counts with MATCHING IDs\n"
        + "6. Sentence positions are relative to the aligned paragraph group (0-based continuous)\n\n"
  
        + "OUTPUT FORMAT - Return EXACTLY this JSON structure (no additional fields):\n"
        + "{\n"
        + "  \"alignments\": [\n"
        + "    {\n"
        + "      \"paragraphAlignment\": {\n"
        + "        \"type\": \"MATCH\",\n"
        + "        \"sourceParagraphPositions\": [0],\n"
        + "        \"targetParagraphPositions\": [0]\n"
        + "      },\n"
        + "      \"sentenceAlignments\": [\n"
        + "        {\n"
        + "          \"type\": \"MATCH\",\n"
        + "          \"sourcePositions\": [0],\n"
        + "          \"targetPositions\": [0]\n"
        + "        },\n"
        + "        {\n"
        + "          \"type\": \"MATCH\",\n"
        + "          \"sourcePositions\": [1],\n"
        + "          \"targetPositions\": [1]\n"
        + "        }\n"
        + "      ]\n"
        + "    },\n"
        + "    {\n"
        + "      \"paragraphAlignment\": {\n"
        + "        \"type\": \"MULTI_MATCH\",\n"
        + "        \"sourceParagraphPositions\": [1, 2],\n"
        + "        \"targetParagraphPositions\": [1]\n"
        + "      },\n"
        + "      \"sentenceAlignments\": [\n"
        + "        {\n"
       + "          \"type\": \"MATCH\",\n"
       + "          \"sourcePositions\": [0],\n"
       + "          \"targetPositions\": [0]\n"
       + "        },\n"
       + "        {\n"
       + "          \"type\": \"MULTI_MATCH\",\n"
       + "          \"sourcePositions\": [1, 2],\n"
       + "          \"targetPositions\": [1]\n"
       + "        }\n"
       + "      ]\n"
       + "    }\n"
       + "  ]\n"
       + "}\n\n"
 
       + "CRITICAL: Use exact field names:\n"
       + "- \"sourceParagraphPositions\" (NOT \"source_paragraphs\")\n"
       + "- \"targetParagraphPositions\" (NOT \"target_paragraphs\")\n"
       + "- \"sourcePositions\" (NOT \"source_positions\")\n"
       + "- \"targetPositions\" (NOT \"target_positions\")\n"
       + "- Values must be INTEGER ARRAYS, not objects with position/id\n\n"
 
       + "Never return anything else. Make all decisions independently.";
 
   private static OpenAIClient client;
 
   public static void init() throws Exception {
     client = OpenAIOkHttpClient.builder()
         .apiKey(Const.ATH_OPENAI_API_KEY)
         .build();
   }
 
   /**
    * Reinserts inline codes into target segments using an LLM to align them with source codes.
    * 
    * <p>
    * This method analyzes source and target text segments and uses an OpenAI model to
    * determine the correct positions for inline codes in the target text based on word/phrase
    * alignment with the source text.
    * </p>
    * 
    * <p>
    * The method guarantees to return a list of the same size as the input layeredSegments,
    * with empty lists for segments that have no codes or when processing fails.
    * </p>
    * 
    * @param modelName       the OpenAI model name to use (e.g., "gpt-5-mini")
    * @param layeredSegments the list of segments containing source and target text with code
    *                        references
    * @return a list of code reference lists, one per input segment. Never null, always the same
    *         size as layeredSegments. Individual lists are never null (empty lists for no codes).
    */
   public static List<List<InlineCodeRef>> reinsertCodes(String modelName,
       List<LayeredSegment> layeredSegments) {
 
     // Validate parameters and return empty results if invalid
     if (Util.isEmpty(modelName) || Util.isEmpty(layeredSegments)) {
       Log.warn(OpenAi.class, "Invalid parameters: modelName='{}', layeredSegments size={}",
           modelName, layeredSegments == null ? "null" : layeredSegments.size());
 
       return layeredSegments == null ? new ArrayList<>()
           : Collections.nCopies(layeredSegments.size(), new ArrayList<>());
     }
 
     // Check if all srcCodes arrays are empty - skip API call if so
     boolean hasAnyCodes = layeredSegments.stream()
         .anyMatch(seg -> seg.getSrcCodes() != null && !seg.getSrcCodes().isEmpty());
 
     if (!hasAnyCodes) {
       Log.debug(OpenAi.class, "All source segments have no codes, skipping API call");
       return Collections.nCopies(layeredSegments.size(), new ArrayList<>());
     }
 
     try {
       String batchJson = JacksonUtil.toJson(layeredSegments, false);
       Log.debug(OpenAi.class, "OpenAI input: {}", batchJson);
 
       ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
           .model(modelName)
           .addSystemMessage(CODES_REINSERTION_PROMPT)
           .addUserMessage(batchJson)
           .build();
 
       // Call the API
       ChatCompletion completion = client.chat().completions().create(params);
 
       // Extract the result (the filled trgCodes JSON array)
       String json = completion.choices().get(0).message().content()
           .orElse("");
 
       if (json.isEmpty()) {
         Log.warn(OpenAi.class, "OpenAI API returned empty content");
         return Collections.nCopies(layeredSegments.size(), new ArrayList<>());
       }
 
       Log.debug(OpenAi.class, "OpenAI output: {}", json);
 
       TypeReference<List<List<InlineCodeRef>>> ref = new TypeReference<>() {
       };
 
       List<List<InlineCodeRef>> trgCodesList = JacksonUtil.fromJson(json, ref);
 
       // Ensure the response is never null and has the correct size
       if (trgCodesList == null) {
         Log.error(OpenAi.class, "Failed to parse OpenAI response, returning empty results");
         return Collections.nCopies(layeredSegments.size(), new ArrayList<>());
       }
 
       // Normalize the result to match expected size
       if (trgCodesList.size() != layeredSegments.size()) {
         Log.warn(OpenAi.class,
             "OpenAI API returned wrong number of results -- expected: {}, actual: {}",
             layeredSegments.size(), trgCodesList.size());
 
         // Pad or trim to match expected size
         List<List<InlineCodeRef>> normalizedList = new ArrayList<>(layeredSegments.size());
 
         for (int i = 0; i < layeredSegments.size(); i++) {
           if (i < trgCodesList.size() && trgCodesList.get(i) != null) {
             normalizedList.add(trgCodesList.get(i));
 
           } else {
             normalizedList.add(new ArrayList<>());
           }
         }
 
         return normalizedList;
       }
 
       // Replace any null lists with empty lists
       for (int i = 0; i < trgCodesList.size(); i++) {
         if (trgCodesList.get(i) == null) {
           trgCodesList.set(i, new ArrayList<>());
         }
       }
 
       return trgCodesList;
 
     } catch (Exception e) {
       Log.error(OpenAi.class, "Exception calling OpenAI API: {}", e.getMessage(), e);
 
       // Return empty lists for all segments so caller can continue
       return Collections.nCopies(layeredSegments.size(), new ArrayList<>());
     }
   }
 
   /**
    * Performs combined paragraph and sentence alignment using an LLM.
    * 
    * <p>
    * This method sends a request to align both paragraphs and sentences in a single API call.
    * The input contains source and target paragraphs with their segmented sentences. The LLM
    * first aligns paragraphs (which may be split, merged, or reordered between documents),
    * then aligns sentences within each aligned paragraph pair.
    * </p>
    * 
    * @param modelName the OpenAI model name to use (e.g., "gpt-4o-mini")
    * @param input     the alignment input containing source and target paragraphs with segments
    * @return the alignment output with paragraph and sentence alignments
    * @throws Exception if the API call fails or returns invalid data
    */
   public static CombinedAlignmentOutput alignParagraphsAndSentences(String modelName,
       CombinedAlignmentInput input) throws Exception {
 
     // Validate parameters
     if (Util.isEmpty(modelName) || input == null) {
       Log.warn(OpenAi.class, "Invalid parameters: modelName='{}', input={}",
           modelName, input == null ? "null" : "present");
       throw new IllegalArgumentException("modelName and input are required");
     }
 
     String inputJson = JacksonUtil.toJson(input, false);
     Log.debug(OpenAi.class, "LLM alignment input: {}", inputJson);
 
     try {
       ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
           .model(modelName)
           .addSystemMessage(COMBINED_ALIGNMENT_PROMPT)
           .addUserMessage(inputJson)
           .build();
 
       ChatCompletion completion = client.chat().completions().create(params);
       String outputJson = completion.choices().get(0).message().content().orElse("");
 
       if (outputJson.isEmpty()) {
         Log.warn(OpenAi.class, "OpenAI API returned empty content");
         throw new Exception("OpenAI API returned empty content");
       }
 
       Log.debug(OpenAi.class, "LLM alignment output: {}", outputJson);
 
       CombinedAlignmentOutput output = JacksonUtil.fromJson(outputJson,
           CombinedAlignmentOutput.class);
 
       if (output == null || output.alignments == null) {
         Log.error(OpenAi.class, "Failed to parse OpenAI response");
         throw new Exception("Failed to parse OpenAI response");
       }
 
       return output;
 
     } catch (Exception e) {
       Log.error(OpenAi.class, "Exception calling OpenAI API: {}", e.getMessage(), e);
       throw e;
     }
   }
 
   public static void done() throws Exception {
     client.close();
   }
 
   public static OpenAIClient getClient() {
     return client;
   }
 }