import { buildSystemPrompt } from "../config/prompt.js";
import { createAdapter } from "./llm-adapters/factory.js";
import { BATCH_CONFIG } from "../config/batch.js";
import type { Language, Pos, EnrichedSense } from "./merge-enriched-data.js";
import { LANG_MAP, POS_MAP, ALL_LANGUAGES } from "../config/constants.js";

interface LlmResponse {
  content: string;
  promptTokens: number;
  completionTokens: number;
  totalTokens: number;
  promptTimeMs: number | null;
  completionTimeMs: number | null;
  totalTimeMs: number;
}

export interface EnrichmentResult {
  results: Map<string, EnrichedSense[]>;
  metrics: {
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
    promptTimeMs: number | null;
    completionTimeMs: number | null;
    totalTimeMs: number;
  };
}

/**
 * Calls the LLM with the enrichment prompt.
 * Returns the response content and timing metrics.
 */
async function callLlm(
  words: string[],
  rawLanguage: string,
  rawPos: string,
): Promise<LlmResponse> {
  const adapter = createAdapter();
  const sourceCode = LANG_MAP[rawLanguage] || rawLanguage;
  const targetLanguages = ALL_LANGUAGES.filter((lang) => lang !== sourceCode);
  const prompt = buildSystemPrompt(rawLanguage, rawPos, targetLanguages);
  return adapter.call(words, prompt);
}

/**
 * Strips markdown code blocks and extracts the JSON object from raw LLM output.
 * Throws if no valid JSON object braces are found.
 */
function sanitizeLlmOutput(raw: string): string {
  const cleaned = raw.replace(/```json\s*/g, "").replace(/```\s*$/g, "");
  const start = cleaned.indexOf("{");
  const end = cleaned.lastIndexOf("}");
  if (start === -1 || end === -1) {
    throw new Error("No JSON object found in LLM output");
  }
  return cleaned.slice(start, end + 1);
}

function validateSense(item: unknown, word: string, index: number): void {
  if (typeof item !== "object" || item === null || Array.isArray(item)) {
    throw new Error(`Sense ${index} for "${word}" is not an object`);
  }

  const sense = item as Record<string, unknown>;

  if (typeof sense["sense"] !== "string" || !sense["sense"]) {
    throw new Error(`Sense ${index} for "${word}": missing or invalid "sense"`);
  }
  if (typeof sense["example"] !== "string" || !sense["example"]) {
    throw new Error(
      `Sense ${index} for "${word}": missing or invalid "example"`,
    );
  }
  if (
    !["easy", "medium", "hard"].includes(sense["difficulty_level"] as string)
  ) {
    throw new Error(`Sense ${index} for "${word}": invalid "difficulty_level"`);
  }
  if (
    typeof sense["translations"] !== "object" ||
    sense["translations"] === null
  ) {
    throw new Error(`Sense ${index} for "${word}": missing "translations"`);
  }

  const trans = sense["translations"] as Record<string, unknown>;
  for (const lang of ["de", "it", "es", "fr"]) {
    if (!Array.isArray(trans[lang])) {
      throw new Error(
        `Sense ${index} for "${word}": missing or invalid "${lang}" translations`,
      );
    }
    for (let j = 0; j < (trans[lang] as unknown[]).length; j++) {
      const t = (trans[lang] as unknown[])[j] as Record<string, unknown>;
      if (typeof t["word"] !== "string" || !t["word"]) {
        throw new Error(
          `Sense ${index} for "${word}": ${lang}[${j}] missing "word"`,
        );
      }
      if (
        !["masculine", "feminine", "neuter", null].includes(
          t["gender"] as string | null,
        )
      ) {
        throw new Error(
          `Sense ${index} for "${word}": ${lang}[${j}] invalid "gender"`,
        );
      }
    }
  }
}

/**
 * Parses the LLM response string into a JavaScript object.
 * Throws if the response is not valid JSON or not an object with expected keys.
 */
export function parseLlmResponse(
  rawJson: string,
  expectedWords: string[],
): Record<string, unknown> {
  let parsed: unknown;

  try {
    const sanitized = sanitizeLlmOutput(rawJson);
    parsed = JSON.parse(sanitized);
  } catch (error: unknown) {
    throw new Error(`Failed to parse LLM output as JSON: ${rawJson}`, {
      cause: error,
    });
  }

  if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
    throw new Error("LLM output is not a JSON object");
  }

  const obj = parsed as Record<string, unknown>;

  for (const word of expectedWords) {
    if (!(word in obj)) {
      throw new Error(`Missing key in LLM output: "${word}"`);
    }
    if (!Array.isArray(obj[word]) || (obj[word] as unknown[]).length === 0) {
      throw new Error(`LLM output for "${word}" is not a non-empty array`);
    }

    // Validate each sense in the array
    const senses = obj[word] as unknown[];
    for (let i = 0; i < senses.length; i++) {
      validateSense(senses[i], word, i);
    }
  }

  return obj;
}

/**
 * Takes parsed LLM output and builds final enriched objects with composite IDs.
 */
export function buildEnrichedData(
  parsed: Record<string, unknown>,
  rawLanguage: string,
  rawPos: string,
): Map<string, EnrichedSense[]> {
  const language = (LANG_MAP[rawLanguage] || rawLanguage) as Language;
  const pos = (POS_MAP[rawPos] || rawPos) as Pos;
  const results = new Map<string, EnrichedSense[]>();

  for (const [word, sensesArray] of Object.entries(parsed)) {
    const senses = (sensesArray as unknown[]).map((item, index) => {
      const sense = item as Omit<
        EnrichedSense,
        "id" | "word" | "language" | "pos"
      >;

      return {
        id: `${word}:${language}:${pos}:${index}`,
        word,
        language,
        pos,
        ...sense,
      } as EnrichedSense;
    });

    results.set(word, senses);
  }

  return results;
}

/**
 * Enriches a batch of words by calling the LLM, parsing the response, and building final data.
 */
export async function enrichWord(
  words: string[],
  rawLanguage: string,
  rawPos: string,
): Promise<EnrichmentResult> {
  const llmResponse = await callLlm(words, rawLanguage, rawPos);
  const parsed = parseLlmResponse(llmResponse.content, words);
  const results = buildEnrichedData(parsed, rawLanguage, rawPos);

  return {
    results,
    metrics: {
      promptTokens: llmResponse.promptTokens,
      completionTokens: llmResponse.completionTokens,
      totalTokens: llmResponse.totalTokens,
      promptTimeMs: llmResponse.promptTimeMs,
      completionTimeMs: llmResponse.completionTimeMs,
      totalTimeMs: llmResponse.totalTimeMs,
    },
  };
}

/**
 * Enriches a batch of words with retry and split-on-failure logic.
 * Retries up to BATCH_CONFIG.maxRetries times, then splits batch in half and retries each half.
 * Continues splitting until batch size is 1, then throws if still failing.
 */
export async function enrichWordWithRetry(
  words: string[],
  rawLanguage: string,
  rawPos: string,
  attempt: number = 1,
): Promise<EnrichmentResult> {
  try {
    return await enrichWord(words, rawLanguage, rawPos);
  } catch (error: unknown) {
    const errorMessage = error instanceof Error ? error.message : String(error);

    if (words.length === 1) {
      throw new Error(
        `Failed to enrich word "${words[0]}" after ${attempt} attempts: ${errorMessage}`,
        { cause: error },
      );
    }

    if (attempt < BATCH_CONFIG.maxRetries) {
      console.log(
        `  Retry ${attempt}/${BATCH_CONFIG.maxRetries} for batch [${words.join(", ")}]: ${errorMessage}`,
      );
      return enrichWordWithRetry(words, rawLanguage, rawPos, attempt + 1);
    }

    // Max retries reached — split and retry
    console.log(
      `  Splitting batch [${words.join(", ")}] after ${BATCH_CONFIG.maxRetries} failed attempts`,
    );

    const half = Math.ceil(words.length / 2);
    const left = words.slice(0, half);
    const right = words.slice(half);

    const leftResult = await enrichWordWithRetry(left, rawLanguage, rawPos, 1);
    const rightResult = await enrichWordWithRetry(
      right,
      rawLanguage,
      rawPos,
      1,
    );

    // Merge results
    const merged = new Map([...leftResult.results, ...rightResult.results]);
    const mergedMetrics = {
      promptTokens:
        leftResult.metrics.promptTokens + rightResult.metrics.promptTokens,
      completionTokens:
        leftResult.metrics.completionTokens +
        rightResult.metrics.completionTokens,
      totalTokens:
        leftResult.metrics.totalTokens + rightResult.metrics.totalTokens,
      promptTimeMs:
        (leftResult.metrics.promptTimeMs ?? 0) +
        (rightResult.metrics.promptTimeMs ?? 0),
      completionTimeMs:
        (leftResult.metrics.completionTimeMs ?? 0) +
        (rightResult.metrics.completionTimeMs ?? 0),
      totalTimeMs:
        leftResult.metrics.totalTimeMs + rightResult.metrics.totalTimeMs,
    };

    return { results: merged, metrics: mergedMetrics };
  }
}