import { normalizeText } from './text_preprocessing.js';

let generator = null;
let loadingPromise = null;
let precomputedMap = null;

const RULES = [
  [/\b[Ss]he go\b/g, (match) => `${match[0]}es`],
  [/\b[Hh]e go\b/g, (match) => `${match[0]}es`],
  [/\b[Ee]very (\w+) have\b/g, 'Every $1 has'],
  [/\b([Tt]he quality results) is\b/g, '$1 are'],
  [/\bdiscussed about\b/gi, 'discussed'],
  [/\bequipments\b/gi, 'equipment'],
  [/\bpeople is\b/gi, 'people are'],
  [/\bthere is many\b/gi, 'there are many'],
  [/\bwe was\b/gi, 'we were'],
  [/\bcan sings\b/gi, 'can sing'],
  [/\bdidn['’]?t went\b/gi, "didn't go"],
  [/\bdidn['’]?t knew\b/gi, "didn't know"],
  [/\bmore better\b/gi, 'better'],
  [/\bmarried with\b/gi, 'married to'],
  [/\blook forward to meet\b/gi, 'look forward to meeting'],
  [/\bI has\b/g, 'I have'],
  [/\bi has\b/g, 'i have'],
  [/\bworks to do\b/gi, 'work to do'],
];

async function loadPrecomputed() {
  if (precomputedMap) return precomputedMap;
  const response = await fetch(new URL('./sample_outputs/precomputed.json', import.meta.url));
  const samples = await response.json();
  precomputedMap = new Map(samples.map((item) => [normalizeText(item.original).toLowerCase(), item]));
  return precomputedMap;
}

function buildPrompt(text) {
  return `Task:\nCorrect the grammar, spelling, punctuation, and wording of the text while preserving the original meaning.\n\nRules:\n- Preserve the original meaning.\n- Do not add unsupported information.\n- Do not remove important details.\n- Keep named entities, numbers, technical terms, and product names unchanged unless clearly incorrect.\n- Return only the corrected text.\n\nInput:\n${text}\n\nCorrected Text:`;
}

export async function loadLiveModel(metadata, onProgress = () => {}) {
  if (!metadata.enable_live_model || metadata.browser_model_id.startsWith('YOUR_')) {
    throw new Error('Live browser model is not configured. The app will use the disclosed fallback mode.');
  }
  if (generator) return generator;
  if (!loadingPromise) {
    loadingPromise = (async () => {
      const { pipeline } = await import('https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1/+esm');
      const options = {
        dtype: metadata.browser_dtype || 'q8',
        device: metadata.preferred_device || 'wasm',
        progress_callback: onProgress,
      };
      generator = await pipeline(metadata.browser_task || 'text2text-generation', metadata.browser_model_id, options);
      return generator;
    })();
  }
  return loadingPromise;
}

function deterministicFallback(text) {
  let corrected = text;
  RULES.forEach(([pattern, replacement]) => { corrected = corrected.replace(pattern, replacement); });
  corrected = corrected.replace(/\bi\b/g, 'I').replace(/\s+([.,!?;:])/g, '$1');
  if (corrected && !/[.!?]$/.test(corrected)) corrected += '.';
  return corrected;
}

export async function correctText(rawText, metadata, onProgress = () => {}) {
  const original = normalizeText(rawText);
  if (!original) throw new Error('Enter text before running correction.');
  const start = performance.now();

  if (metadata.enable_live_model && !metadata.browser_model_id.startsWith('YOUR_')) {
    const pipe = await loadLiveModel(metadata, onProgress);
    const output = await pipe(buildPrompt(original), { max_new_tokens: 256, num_beams: 4, do_sample: false });
    const corrected = output?.[0]?.generated_text?.trim() ?? '';
    return { original, corrected, latencyMs: performance.now() - start, mode: 'live-browser-model', confidenceSource: 'generation output proxy' };
  }

  const examples = await loadPrecomputed();
  const sample = examples.get(original.toLowerCase());
  const corrected = sample?.corrected ?? deterministicFallback(original);
  return {
    original,
    corrected,
    latencyMs: performance.now() - start,
    mode: sample ? 'precomputed-fallback-sample' : 'rule-based-fallback-demo',
    confidenceSource: 'heuristic only',
    disclosure: metadata.fallback_disclosure,
  };
}