import { normalize, tokenAffinity } from "../text.js";

const MAX_EXCERPT_CHARACTERS = 900;
const EXCERPT_WINDOW_LINES = 21;
const TOKEN_AFFINITY = 0.55;
const EXCERPT_CANDIDATES = 64;
const CODE_ALIASES = new Map([
  ["function", ["fn", "func", "def"]],
  ["functions", ["fn", "func", "def"]],
]);

/** @param {string[]} words */
function searchTokens(words) {
  return [...new Set(words.flatMap((word) => [
    word,
    ...(CODE_ALIASES.get(word) ?? []),
  ]))];
}

const DECLARATION = /^\s*(?:(?:pub(?:\([^)]*\))?|export|default|async|static)\s+)*(?:fn|function|class|struct|enum|trait|interface|impl)\b/;

/** @param {string[]} lines */
function declarations(lines) {
  return lines
    .map((line, index) => DECLARATION.test(line) ? index : -1)
    .filter((index) => index >= 0);
}

/** @param {string} text @param {number} line @param {number | undefined} declaration */
function excerptAt(text, line, declaration) {
  const lines = text.split("\n");
  const declarationLine = declaration === undefined ? undefined : declaration + 1;
  const nearbyDeclaration = declarationLine !== undefined &&
    line - declarationLine < EXCERPT_WINDOW_LINES - 3;
  const start = nearbyDeclaration ? declarationLine : Math.max(1, line - 3);
  const end = Math.min(
    lines.length,
    nearbyDeclaration
      ? start + EXCERPT_WINDOW_LINES - 1
      : line + 12,
  );
  const local = lines.slice(start - 1, end).join("\n");
  const content = declarationLine !== undefined && declarationLine < start
    ? `${lines[declarationLine - 1]}\n...\n${local}`
    : local;
  return {
    startLine: declarationLine !== undefined && declarationLine < start
      ? declarationLine
      : start,
    endLine: end,
    content: content.slice(0, MAX_EXCERPT_CHARACTERS),
  };
}

/** @param {string} text @param {string} phrase */
function phraseOccurrences(text, phrase) {
  if (!phrase) return 0;
  let count = 0;
  let index = 0;
  while ((index = text.indexOf(phrase, index)) >= 0) {
    count += 1;
    index += phrase.length;
  }
  return count;
}

/** @param {string} content */
function codeOnly(content) {
  return content
    .replace(/\/\*[\s\S]*?\*\//g, "")
    .split("\n")
    .map((line) => {
      const trimmed = line.trimStart();
      return trimmed === "*" ||
          trimmed.startsWith("* ") ||
          trimmed.startsWith("*/")
        ? ""
        : line.split("//")[0];
    })
    .join("\n");
}

/** @param {string} content */
function computationWeight(content) {
  const code = codeOnly(content);
  const arithmetic = code.includes("+=") ||
    code.includes("-=") ||
    code.includes("*=") ||
    code.includes("/=") ||
    code.includes(" + ") ||
    code.includes(" - ") ||
    code.includes(" * ") ||
    code.includes(" / ");
  return arithmetic ? 2 : code.includes(".compute") ? 1 : 0;
}

/** @param {string} value */
function regexEscape(value) {
  return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}

/** @param {string} line @param {string} query */
function callsQuery(line, query) {
  if (!/^[a-z_][a-z0-9_]*$/i.test(query)) {
    return line.includes("(") && normalize(line).includes(query);
  }
  return new RegExp(`\\b${regexEscape(query)}\\s*\\(`).test(line);
}

/** @param {string} content @param {string} query */
function usageWeight(content, query) {
  return codeOnly(content).split("\n").filter((line) =>
    !DECLARATION.test(line) &&
    callsQuery(line, query)
  ).length;
}

/** @param {string} content @param {string} query */
function computesQueryDirectly(content, query) {
  return codeOnly(content).split("\n").some((line) => {
    return normalize(line).includes(query) && computationWeight(line) > 0;
  });
}

/** @param {string} content @param {string} query */
function containsDirectFormula(content, query) {
  return content.split("\n").some((line) => {
    const assignment = line.match(/^(.+?)(?:\+=|-=|\*=|\/=|=)(.+)$/);
    return assignment &&
      normalize(assignment[1]).includes(query) &&
      /[+\-*/×÷Σ∑]/.test(assignment[2]);
  });
}

/** @param {{ path: string, text: string }[]} files */
export function createSourceSearchIndex(files) {
  /** @type {Map<string, number[]>} */
  const postings = new Map();

  files.forEach((file, fileIndex) => {
    const tokens = new Set(
      normalize(`${file.path}\n${file.text}`).split(" ").filter(Boolean),
    );
    for (const token of tokens) {
      const matches = postings.get(token);
      if (matches) matches.push(fileIndex);
      else postings.set(token, [fileIndex]);
    }
  });

  return {
    files,
    paths: files.map((file) => normalize(file.path)),
    postings,
  };
}

/** @param {string} left @param {string} right */
function relatedToken(left, right) {
  return left === right || tokenAffinity(left, right) >= TOKEN_AFFINITY;
}

/** @param {string} text @param {string[]} tokens @param {number[]} weights @param {string} phrase @param {boolean} preferComputation @param {boolean} preferUsage @param {string} usageQuery */
function bestExcerptLine(
  text,
  tokens,
  weights,
  phrase,
  preferComputation,
  preferUsage,
  usageQuery,
) {
  const lines = text.split("\n");
  const normalized = lines.map((line) => normalize(line));
  const strongest = Math.max(...weights, 1);
  const starts = declarations(lines);
  const scopes = starts.length
    ? starts.map((start, index) => ({
        start,
        end: (starts[index + 1] ?? lines.length) - 1,
      }))
    : [{ start: 0, end: lines.length - 1 }];
  const rankedScopes = scopes.map((scope) => {
    const scopeText = normalized.slice(scope.start, scope.end + 1).join(" ");
    const words = new Set(
      scopeText.split(" ")
        .filter(Boolean),
    );
    let score = tokens.reduce((sum, token, tokenIndex) =>
      [...words].some((word) => relatedToken(word, token))
        ? sum + weights[tokenIndex]
        : sum, 0);
    if (phrase && scopeText.includes(phrase)) {
      score += strongest * 3;
    }
    if (preferComputation) {
      score += computationWeight(
        lines.slice(scope.start, scope.end + 1).join("\n"),
      ) * strongest * 4;
    }
    if (preferUsage) {
      score += usageWeight(
        lines.slice(scope.start, scope.end + 1).join("\n"),
        usageQuery,
      ) * strongest * 4;
    }
    return { ...scope, score };
  }).sort((left, right) => right.score - left.score || left.start - right.start);
  const scope = rankedScopes[0];
  let bestLine = 1;
  let bestScore = 0;

  for (let index = scope.start; index <= scope.end; index += 1) {
    const line = normalized[index];
    const nearby = normalized
      .slice(
        Math.max(0, index - Math.floor(EXCERPT_WINDOW_LINES / 2)),
        index + Math.ceil(EXCERPT_WINDOW_LINES / 2),
      )
      .join(" ");
    const lineWords = new Set(line.split(" ").filter(Boolean));
    const nearbyWords = new Set(nearby.split(" ").filter(Boolean));
    let score = tokens.reduce((sum, token, tokenIndex) => {
      if ([...lineWords].some((word) => relatedToken(word, token))) {
        return sum + weights[tokenIndex];
      }
      return [...nearbyWords].some((word) => relatedToken(word, token))
        ? sum + weights[tokenIndex] * 0.2
        : sum;
    }, 0);
    if (phrase && line.includes(phrase)) score += strongest * 3;
    if (
      preferUsage &&
      !DECLARATION.test(lines[index]) &&
      callsQuery(lines[index], usageQuery)
    ) score += strongest * 4;
    if (!score || score <= bestScore) continue;
    bestScore = score;
    bestLine = index + 1;
  }
  const declaration = starts
    .filter((start) => start <= bestLine - 1)
    .at(-1);
  return { line: bestLine, declaration };
}

/**
 * @param {ReturnType<typeof createSourceSearchIndex>} index
 * @param {string} token
 */
function candidateFiles(index, token) {
  const exact = index.postings.get(token);
  const matches = new Set(exact ?? []);
  if (token.length < 5) return [...matches];

  for (const [candidate, files] of index.postings) {
    if (candidate === token) continue;
    if (!relatedToken(candidate, token)) continue;
    for (const file of files) matches.add(file);
  }
  return [...matches];
}

/**
 * @param {ReturnType<typeof createSourceSearchIndex>} index
 * @param {string} rawQuery
 * @param {string} [pathPrefix]
 * @param {"definition" | "implementation" | "usage" | "availability"} [focus]
 */
export function searchSource(index, rawQuery, pathPrefix = "", focus = undefined) {
  const exactUsageQuery = /^[A-Za-z_][A-Za-z0-9_]*$/.test(rawQuery.trim())
    ? rawQuery.trim()
    : normalize(rawQuery);
  const query = normalize(rawQuery);
  if (!query) throw new Error("Search query is empty");

  const words = query.split(" ");
  const tokens = searchTokens(words);
  const tokenMatches = tokens
    .map((token) => ({ token, files: candidateFiles(index, token) }))
    .filter(({ files }) => files.length)
    .map(({ token, files }) => ({
      token,
      files,
      weight: Math.log2((index.files.length + 1) / (files.length + 1)) + 1,
    }));
  if (!tokenMatches.length) return [];
  /** @type {Map<number, { weights: number[], queryTokens: string[], matched: number }>} */
  const candidates = new Map();
  for (const { token, files, weight } of tokenMatches) {
    for (const fileIndex of files) {
      const candidate = candidates.get(fileIndex) ?? {
        weights: [],
        queryTokens: [],
        matched: 0,
      };
      candidate.weights.push(weight);
      candidate.queryTokens.push(token);
      candidate.matched += 1;
      candidates.set(fileIndex, candidate);
    }
  }

  const matches = [];
  for (const [fileIndex, candidate] of candidates) {
    const file = index.files[fileIndex];
    if (pathPrefix && !file.path.startsWith(pathPrefix)) continue;

    const normalized = normalize(file.text);
    const exact = normalized.includes(query);
    const normalizedPath = index.paths[fileIndex];
    const pathTokens = normalizedPath.split(" ");
    const pathScore = tokenMatches.reduce((score, { token, weight }) =>
      pathTokens.some((pathToken) => relatedToken(pathToken, token))
        ? score + weight
        : score, 0);
    const semanticScore = [...candidate.weights]
      .sort((left, right) => right - left)
      .slice(0, 3)
      .reduce((sum, weight) => sum + weight, 0);
    matches.push({
      fileIndex,
      path: file.path,
      score: semanticScore +
        pathScore * 2 +
        (exact ? 30 : 0) +
        (normalizedPath.includes(query) ? 15 : 0),
      matched: candidate.matched,
      queryTerms: tokens.length,
      queryTokens: candidate.queryTokens,
      weights: candidate.weights,
      pathTokens,
    });
  }

  const ranked = matches.sort((left, right) =>
      right.score - left.score ||
      right.matched - left.matched ||
      left.path.localeCompare(right.path)
    )
    .slice(0, EXCERPT_CANDIDATES)
    .map((match) => {
      const file = index.files[match.fileIndex];
      const { line, declaration } = bestExcerptLine(
        file.text,
        match.queryTokens,
        match.weights,
        query,
        focus === "implementation",
        focus === "usage",
        exactUsageQuery,
      );
      const excerpt = excerptAt(file.text, line, declaration);
      const localPhraseOccurrences = phraseOccurrences(
        normalize(excerpt.content),
        query,
      );
      const declarationText = declaration === undefined
        ? ""
        : normalize(file.text.split("\n")[declaration]);
      const definitionScore = focus === "definition" &&
          (
            declarationText.includes(query) ||
            excerpt.content.split("\n").some((line) =>
              DECLARATION.test(line) && normalize(line).includes(query)
            )
          )
        ? 60
        : 0;
      const computesDirectly = focus === "implementation" &&
        computesQueryDirectly(excerpt.content, query);
      const implementationScore = computesDirectly
        ? computationWeight(excerpt.content) * 30
        : 0;
      const directImplementationScore = computesDirectly ? 40 : 0;
      const formulaScore = containsDirectFormula(excerpt.content, query)
        ? 80
        : 0;
      const usageScore = focus === "usage"
        ? usageWeight(excerpt.content, exactUsageQuery) * 40
        : 0;
      return {
        ...match,
        score: match.score +
          localPhraseOccurrences * 10 +
          definitionScore +
          implementationScore +
          directImplementationScore +
          formulaScore +
          usageScore,
        phraseOccurrences: localPhraseOccurrences,
        ...excerpt,
      };
    })
    .sort((left, right) =>
      right.score - left.score ||
      right.phraseOccurrences - left.phraseOccurrences ||
      right.matched - left.matched ||
      left.path.localeCompare(right.path)
    );
  const diversified = ranked.slice(0, 8);
  const salient = [...tokenMatches]
    .sort((left, right) => right.weight - left.weight)
    .slice(0, 5);
  for (const { token } of salient) {
    const candidate = ranked
      .filter((match) => match.queryTokens.includes(token))
      .sort((left, right) =>
        Number(right.pathTokens.some((pathToken) => relatedToken(pathToken, token))) -
          Number(left.pathTokens.some((pathToken) => relatedToken(pathToken, token))) ||
        right.score - left.score
      )[0];
    if (candidate) diversified.push(candidate);
  }

  return [...new Map(diversified.map((match) => [match.path, match])).values()]
    .slice(0, 8)
    .map(({ fileIndex, queryTokens, weights, pathTokens, ...match }) => match);
}