// Initialize the app container.
const app = document.getElementById('app');

// Add retry count with 600 second timeout (checking every 2 seconds = 300 attempts).
let retryCount = 0;
const MAX_RETRIES = 300; // 600 seconds / 2 second interval = 300 attempts

function checkServicesStatus() {
  fetch('/inference/api/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json'
    },
    body: JSON.stringify({
      model: 'gemma-2b',
      messages: [{role: 'user', content: 'hi'}]
    })
  })
    .then(response => {
      if (response.ok) {
        window.servicesReady = true;
        const translationControls = document.getElementById('translationControls');
        const statusIndicator = document.getElementById('statusIndicator');
        const loadingOverlay = document.getElementById('loadingOverlay');
        if (translationControls) {
          translationControls.style.display = 'block';
        }
        if (statusIndicator) {
          statusIndicator.style.display = 'none';
        }
        if (loadingOverlay) {
          loadingOverlay.classList.add('hidden');
        }
        return true;
      }
      throw new Error('Services not ready');
    })
    .catch(error => {
      retryCount++;
      const statusIndicator = document.getElementById('statusIndicator');
      if (statusIndicator) {
        if (retryCount >= MAX_RETRIES) {
          statusIndicator.innerHTML = '❌ Failed to initialize AI services. Try restarting the space.';
          statusIndicator.style.color = '#dc3545';
          clearInterval(statusInterval);
        } else {
          statusIndicator.innerHTML = `<span class="spinner"></span> Initializing AI services...`;
        }
      }
      console.log('Waiting for services...', error);
      return false;
    });
}

// Check status every 2 seconds until ready or max retries reached.
const statusInterval = setInterval(() => {
  if (window.servicesReady || retryCount >= MAX_RETRIES) {
    clearInterval(statusInterval);
  } else {
    checkServicesStatus();
  }
}, 2000);

// Load Document Authoring SDK.
const script = document.createElement('script');
script.src = 'https://document-authoring-cdn.pspdfkit.com/releases/document-authoring-1.0.26-umd.js';
script.crossOrigin = true;
document.head.appendChild(script);

script.onload = async () => {
  // Fetch license key from nginx endpoint.
  const licenseKeyResponse = await fetch('/api/license-key');
  if (!licenseKeyResponse.ok) {
    throw new Error('Failed to fetch license key');
  }
  const { licenseKey } = await licenseKeyResponse.json();
  
  const docAuthSystem = await DocAuth.createDocAuthSystem({licenseKey});
  
  const editor = await docAuthSystem.createEditor(
    document.getElementById("editor"),
    {
      document: await docAuthSystem.importDOCX(fetch('./Sample.docx')),
    }
  );

  // Helper function to extract text runs with their paths.
  function flattenTextRuns(obj, path = [], result = []) {
    for (let key in obj) {
      const newPath = [...path, key];
      if (typeof obj[key] === "string" && key === "text" && obj[key].length > 0) {
        result.push({ text: obj[key], path: newPath });
      } else if (typeof obj[key] === "object") {
        flattenTextRuns(obj[key], newPath, result);
      }
    }
    return result;
  }

  // Helper function to group adjacent text runs
  function groupAdjacentRuns(textRuns, maxChunkSize = 1000) {
    const groups = [];
    let currentGroup = [];
    let currentSize = 0;

    for (let run of textRuns) {
      if (currentSize + run.text.length > maxChunkSize && currentGroup.length > 0) {
        groups.push([...currentGroup]);
        currentGroup = [];
        currentSize = 0;
      }
      currentGroup.push(run);
      currentSize += run.text.length;
    }
    
    if (currentGroup.length > 0) {
      groups.push(currentGroup);
    }
    
    return groups;
  }

  async function translate(content, targetLang, sourceLang = 'English') {
    try {
      const response = await fetch('/inference/api/v1/chat/completions', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          messages: [
            {
              role: "system",
              content: `CRITICAL INSTRUCTION: The word "Nutrient" is a company name and must stay EXACTLY as "Nutrient" in the translation - never translate it to "Nutriente" or any other word.

Translate the following text from ${sourceLang} to ${targetLang}. 
Each segment is marked with ||| delimiters.

Additional rules:
1. Do NOT translate proper nouns (names of people, companies, products, places)
2. Preserve original capitalization and punctuation
3. Add spaces at the start or end of segments only if they exist in the original text

Example:
EN: "Companies use Nutrient to..."
${targetLang}: "Las empresas usan Nutrient para..."`
            },
            {
              role: "user",
              content: content
            }
          ],
          model: "gemma-2b",
          stream: false
        })
      });

      if (!response.ok) {
        throw new Error('Translation request failed');
      }

      const data = await response.json();
      return data.choices[0].message.content;
    } catch (error) {
      console.error("Translation error:", error);
      throw new Error("Failed to generate translation");
    }
  }

  // Function to update nested object value by path.
  function setNestedValue(obj, path, value) {
    let current = obj;
    for (let i = 0; i < path.length - 1; i++) {
      current = current[path[i]];
    }
    current[path[path.length - 1]] = value;
  }

  async function translateRecursive(docJson, targetLang, sourceLang = 'English') {
    const textRuns = flattenTextRuns(docJson);
    const groups = groupAdjacentRuns(textRuns);
    
    for (let group of groups) {
      const content = group.map(run => `|||${run.text}|||`).join('\n');
      const translatedContent = await translate(content, targetLang, sourceLang);
      
      const translations = translatedContent
        .split('|||')
        .filter(t => t.trim())
        .map(t => t.trim());
      
      group.forEach((run, index) => {
        if (translations[index]) {
          const originalText = run.text;
          let translatedText = translations[index];
          
          if (originalText.startsWith(' ')) {
            translatedText = ' ' + translatedText;
          }
          if (originalText.endsWith(' ')) {
            translatedText = translatedText + ' ';
          }
          
          setNestedValue(docJson, run.path, translatedText);
        }
      });
    }
  }

  // Add language selection handlers.
  const sourceLanguageSelect = document.getElementById('sourceLanguageSelect');
  const targetLanguageSelect = document.getElementById('targetLanguageSelect');

  // Update target language options when source language changes.
  sourceLanguageSelect.addEventListener('change', () => {
    const selectedSource = sourceLanguageSelect.value;
    const languages = ['English', 'Spanish', 'French', 'German'];
    
    // Clear existing options.
    targetLanguageSelect.innerHTML = '';
    
    // Add all languages except the selected source language.
    languages.forEach(lang => {
      if (lang !== selectedSource) {
        const option = document.createElement('option');
        option.value = lang;
        option.textContent = lang;
        targetLanguageSelect.appendChild(option);
      }
    });
  });

  const translateButton = document.getElementById('translateButton');
  translateButton.addEventListener('click', async () => {
    try {
      const sourceLanguage = sourceLanguageSelect.value;
      const targetLanguage = targetLanguageSelect.value;

      translateButton.disabled = true;
      translateButton.innerHTML = '<span class="spinner"></span> Translating...';

      const jsonDoc = await editor.currentDocument().saveDocument();
      await translateRecursive(jsonDoc.container.document, targetLanguage, sourceLanguage);
      const translatedDoc = await docAuthSystem.loadDocument(jsonDoc);
      editor.setCurrentDocument(translatedDoc);
    } catch (error) {
      alert('Translation failed: ' + error.message);
    } finally {
      translateButton.disabled = false;
      translateButton.innerHTML = 'Translate Document';
    }
  });
};