import { LLM_TIMEOUT_MS } from "@/constants";
import { TimeoutError } from "@/error";
import { logError, logInfo, logWarn } from "@/logger";
import { extractTagsFromQuery } from "@/search/v3/utils/tagUtils";
import { withSuppressedTokenWarnings, withTimeout } from "@/utils";
import { BaseChatModel } from "@langchain/core/language_models/chat_models";

export interface QueryExpanderOptions {
  maxVariants?: number;
  timeout?: number;
  cacheSize?: number;
  getChatModel?: () => Promise<BaseChatModel | null>;
}

export interface ExpandedQuery {
  queries: string[]; // Original query + expanded variants
  salientTerms: string[]; // Important terms extracted ONLY from original query (used for scoring)
  originalQuery: string; // The original user query
  expandedQueries: string[]; // Only the expanded variants (not including original)
}

/**
 * Expands search queries using LLM to generate alternative phrasings
 * and extracts salient terms for improved search relevance.
 */
export class QueryExpander {
  private cache = new Map<string, ExpandedQuery>();
  private readonly config;

  private static readonly PROMPT_TEMPLATE = `Analyze this search query and provide:
1. SALIENT TERMS from the original query (for ranking)
2. Alternative queries and related terms (for finding more results)

Query: "{query}"

Instructions:

SALIENT TERMS (for ranking - ONLY from original query):
- Extract meaningful words FROM THE ORIGINAL QUERY ONLY
- Include: nouns, proper nouns, technical terms, domain concepts
- Exclude: action verbs (find, search, get), pronouns (my, your), articles (the, a), prepositions (in, on, for), conjunctions (and, or)
- These terms determine search ranking - must be from original query

ALTERNATIVE QUERIES (for recall - specific alternative phrasings):
- Alternative phrasings of the query (specific enough to be useful for search)

Example: "find my piano notes"
- Salient (from original): piano, notes
- Queries: "piano lesson notes", "piano practice sheets"

Example: "查找我的学习笔记"
- Salient (from original): 学习, 笔记
- Queries: "个人笔记文档"

Format:
<salient>
<term>word_from_original_query</term>
</salient>
<queries>
<query>alternative query</query>
</queries>`;

  constructor(private readonly options: QueryExpanderOptions = {}) {
    this.config = {
      maxVariants: options.maxVariants ?? 2,
      timeout: options.timeout ?? LLM_TIMEOUT_MS,
      cacheSize: options.cacheSize ?? 100,
      minTermLength: 2,
    };
  }

  /**
   * Expands a search query into multiple variants and extracts salient terms.
   * Uses caching to avoid redundant LLM calls.
   * @param query - The original search query
   * @returns Expanded queries and salient terms extracted from the original query
   */
  async expand(query: string): Promise<ExpandedQuery> {
    // Check if query is valid
    if (!query?.trim()) {
      return {
        queries: [],
        salientTerms: [],
        originalQuery: "",
        expandedQueries: [],
      };
    }

    // Check cache first (and update LRU position)
    const cached = this.cache.get(query);
    if (cached) {
      // Move to end (most recently used) for proper LRU
      this.cache.delete(query);
      this.cache.set(query, cached);
      logInfo(`QueryExpander: Using cached expansion for "${query}"`);
      return cached;
    }

    try {
      // Expand with timeout protection
      const expanded = await this.expandWithTimeout(query);

      // Cache the result
      this.cacheResult(query, expanded);

      return expanded;
    } catch (error) {
      logWarn(`QueryExpander: Failed to expand query "${query}":`, error);
      // Fallback: extract terms from original query only
      return this.fallbackExpansion(query);
    }
  }

  /**
   * Expands query with timeout protection to prevent hanging on slow LLM responses.
   * @param query - The query to expand
   * @returns Expanded query or fallback if timeout is reached
   */
  private async expandWithTimeout(query: string): Promise<ExpandedQuery> {
    try {
      return await withTimeout(
        (signal) => this.expandWithLLM(query, signal),
        this.config.timeout,
        "Query expansion"
      );
    } catch (error: any) {
      if (error instanceof TimeoutError) {
        logInfo(`QueryExpander: Timeout reached for "${query}"`);
        return this.fallbackExpansion(query);
      }
      throw error;
    }
  }

  /**
   * Performs the actual LLM call to expand the query.
   * @param query - The query to expand
   * @param signal - Optional abort signal for request cancellation
   * @returns Expanded queries and extracted salient terms
   */
  private async expandWithLLM(query: string, signal?: AbortSignal): Promise<ExpandedQuery> {
    try {
      if (!this.options.getChatModel) {
        logInfo("QueryExpander: No chat model getter provided");
        return this.fallbackExpansion(query);
      }

      const model = await this.options.getChatModel();
      if (!model) {
        logInfo("QueryExpander: No chat model available");
        return this.fallbackExpansion(query);
      }

      const prompt = QueryExpander.PROMPT_TEMPLATE.replace(
        "{count}",
        this.config.maxVariants.toString()
      ).replace("{query}", query);

      // Invoke model with token warnings suppressed and abort signal
      const response = await withSuppressedTokenWarnings(async () => {
        return await model.invoke(prompt, signal ? { signal } : undefined);
      });

      if (!response) {
        return this.fallbackExpansion(query);
      }

      const content = this.extractContent(response);

      if (!content) {
        return this.fallbackExpansion(query);
      }

      // Parse response using XML tags
      const parsed = this.parseXMLResponse(content, query);

      logInfo(
        `QueryExpander: Expanded "${query}" to ${parsed.queries.length} queries and ${parsed.salientTerms.length} terms`
      );
      return parsed;
    } catch (error) {
      logError("QueryExpander: LLM expansion failed:", error);
      return this.fallbackExpansion(query);
    }
  }

  /**
   * Extracts text content from various LLM response formats.
   * @param response - The LLM response object or string
   * @returns The extracted text content or null if empty
   */
  private extractContent(response: any): string | null {
    // Elegant extraction with nullish coalescing
    return typeof response === "string"
      ? response
      : String(response?.content ?? response?.text ?? "").trim() || null;
  }

  /**
   * Extracts salient terms from the original query only (used for scoring).
   * @param originalQuery - The original user query
   * @returns Array of valid terms extracted from the original query
   */
  private extractSalientTermsFromOriginal(originalQuery: string): string[] {
    const baseTerms = this.extractTermsFromQueries([originalQuery]);
    const tagTerms = extractTagsFromQuery(originalQuery);
    return this.combineBaseAndTagTerms(baseTerms, tagTerms, originalQuery);
  }

  /**
   * Parses XML-formatted LLM response to extract queries and terms.
   * Falls back to legacy format if XML tags are not found.
   * @param content - The LLM response content
   * @param originalQuery - The original query for fallback term extraction
   * @returns Parsed expanded queries and salient terms
   */
  private parseXMLResponse(content: string, originalQuery: string): ExpandedQuery {
    const queries: string[] = [originalQuery]; // Always include original
    const salientFromLLM = new Set<string>(); // Salient terms from original query (for ranking)

    // Extract queries from XML tags
    const queryRegex = /<query>(.*?)<\/query>/g;
    let queryMatch;
    while ((queryMatch = queryRegex.exec(content)) !== null) {
      const query = queryMatch[1]?.trim();
      if (query && query !== originalQuery && queries.length <= this.config.maxVariants) {
        queries.push(query);
      }
    }

    // Extract SALIENT terms (from original query, for ranking)
    const salientSection = content.match(/<salient>([\s\S]*?)<\/salient>/);
    if (salientSection) {
      const termRegex = /<term>(.*?)<\/term>/g;
      let termMatch;
      while ((termMatch = termRegex.exec(salientSection[1])) !== null) {
        const term = termMatch[1]?.trim().toLowerCase();
        if (term && this.isValidTerm(term)) {
          salientFromLLM.add(term);
        }
      }
    }

    // Check if content has any XML structure (queries or terms)
    const hasXMLContent = queries.length > 1 || salientFromLLM.size > 0 || /<term>/.test(content);

    // If no XML tags found at all, fall back to extracting from original query
    if (!hasXMLContent) {
      return this.fallbackExpansion(originalQuery);
    }

    // Salient terms: from LLM's <salient> section, or fallback to extracting from original query
    const tagTerms = extractTagsFromQuery(originalQuery);
    const salientTerms =
      salientFromLLM.size > 0
        ? Array.from(new Set([...salientFromLLM, ...tagTerms]))
        : this.extractSalientTermsFromOriginal(originalQuery);

    const expandedQueries = queries.slice(1);
    return {
      queries: queries.slice(0, this.config.maxVariants + 1),
      salientTerms: salientTerms, // From original query only (for ranking)
      originalQuery: originalQuery,
      expandedQueries: expandedQueries.slice(0, this.config.maxVariants),
    };
  }

  /**
   * Provides a fallback expansion when LLM is unavailable or fails.
   * Extracts terms directly from the original query and generates fuzzy variants.
   * @param query - The original query
   * @returns Fallback expansion with original query, fuzzy variants, and extracted terms
   */
  private fallbackExpansion(query: string): ExpandedQuery {
    const baseTerms = this.extractTermsFromQueries([query]);
    const tagTerms = extractTagsFromQuery(query);
    const terms = this.combineBaseAndTagTerms(baseTerms, tagTerms, query);

    return {
      queries: [query],
      salientTerms: terms,
      originalQuery: query,
      expandedQueries: [],
    };
  }

  /**
   * Extracts individual terms from queries by splitting and filtering.
   * Handles hyphenated words by extracting both compound and component terms.
   * @param queries - Array of queries to extract terms from
   * @returns Array of unique valid terms
   */
  private extractTermsFromQueries(queries: string[]): string[] {
    const terms = new Set<string>();

    for (const query of queries) {
      // Split on common delimiters and spaces
      const words = query
        .toLowerCase()
        .replace(/[^\w\s-]/g, " ") // Replace punctuation with spaces
        .split(/\s+/);

      for (const word of words) {
        if (this.isValidTerm(word)) {
          terms.add(word);

          // Also add compound terms split by hyphens
          if (word.includes("-")) {
            word.split("-").forEach((part) => {
              if (this.isValidTerm(part)) {
                terms.add(part);
              }
            });
          }
        }
      }
    }

    return Array.from(terms);
  }

  /**
   * Validates if a term should be included in salient terms.
   * Filters out terms that are too short or contain only special characters.
   * Note: Stopword filtering is handled by the LLM prompt, not here.
   * @param term - The term to validate
   * @returns true if the term is valid for inclusion
   */
  private isValidTerm(term: string): boolean {
    if (term.length < this.config.minTermLength) {
      return false;
    }

    if (term.startsWith("#")) {
      try {
        return /^#[\p{L}\p{N}_/-]+$/u.test(term);
      } catch {
        return /^#[A-Za-z0-9_/-]+$/.test(term);
      }
    }

    try {
      return /^[\p{L}\p{N}_-]+$/u.test(term);
    } catch {
      return /^[A-Za-z0-9_-]+$/.test(term);
    }
  }

  /**
   * Merges base terms with tag-prefixed terms while preserving standalone terms present in the query.
   * Removes tag bodies only when they originate exclusively from tag tokens.
   *
   * @param baseTerms - Terms extracted from the raw query (sans tag awareness)
   * @param tagTerms - Hash-prefixed tag tokens extracted from the query
   * @param originalQuery - The original user-supplied query
   * @returns Array of unique salient terms preserving tag intent
   */
  private combineBaseAndTagTerms(
    baseTerms: string[],
    tagTerms: string[],
    originalQuery: string
  ): string[] {
    const combined = new Set<string>([...baseTerms, ...tagTerms]);

    if (tagTerms.length === 0) {
      return Array.from(combined);
    }

    const standaloneTerms = this.collectStandaloneTerms(originalQuery);

    for (const tag of tagTerms) {
      const withoutHash = tag.slice(1);
      if (withoutHash.length > 0 && !standaloneTerms.has(withoutHash)) {
        combined.delete(withoutHash);
      }
    }

    return Array.from(combined);
  }

  /**
   * Collects lowercase standalone terms from the original query that do not belong to tag tokens.
   * These terms originate from user text outside of hash-prefixed tags.
   *
   * @param originalQuery - Raw user query containing potential tag and non-tag text
   * @returns Set of standalone terms detected outside tag spans
   */
  private collectStandaloneTerms(originalQuery: string): Set<string> {
    const standaloneTerms = new Set<string>();

    if (!originalQuery) {
      return standaloneTerms;
    }

    const normalizedQuery = originalQuery.toLowerCase();
    const tagRanges = this.findTagRanges(normalizedQuery);

    const wordPatterns = [/[\p{L}\p{N}_-]+/gu, /[a-z0-9_-]+/g];

    for (const pattern of wordPatterns) {
      try {
        for (const match of normalizedQuery.matchAll(pattern)) {
          if (match.index === undefined) {
            continue;
          }

          const start = match.index;
          const end = start + match[0].length;
          const insideTag = tagRanges.some(
            ({ start: tagStart, end: tagEnd }) => start >= tagStart && end <= tagEnd
          );

          if (insideTag) {
            continue;
          }

          const candidate = match[0];
          if (this.isValidTerm(candidate) && !candidate.startsWith("#")) {
            standaloneTerms.add(candidate);

            if (candidate.includes("-")) {
              candidate.split("-").forEach((part) => {
                if (this.isValidTerm(part) && !part.startsWith("#")) {
                  standaloneTerms.add(part);
                }
              });
            }
          }
        }
        break;
      } catch {
        continue;
      }
    }

    return standaloneTerms;
  }

  /**
   * Finds the start and end indices for every tag token inside the query string.
   * Ranges include the '#' prefix so downstream checks can exclude tag spans precisely.
   *
   * @param normalizedQuery - Lowercase query used for regex scanning
   * @returns Array of inclusive-exclusive index ranges for each tag occurrence
   */
  private findTagRanges(normalizedQuery: string): Array<{ start: number; end: number }> {
    const ranges: Array<{ start: number; end: number }> = [];
    const tagPatterns = [/#[\p{L}\p{N}_/-]+/gu, /#[a-z0-9_/-]+/g];

    for (const pattern of tagPatterns) {
      try {
        for (const match of normalizedQuery.matchAll(pattern)) {
          if (match.index === undefined) {
            continue;
          }

          ranges.push({
            start: match.index,
            end: match.index + match[0].length,
          });
        }
        break;
      } catch {
        continue;
      }
    }

    return ranges;
  }

  /**
   * Caches an expansion result with simple LRU eviction.
   * @param query - The original query as cache key
   * @param expanded - The expansion result to cache
   */
  private cacheResult(query: string, expanded: ExpandedQuery): void {
    // Map maintains insertion order for simple LRU
    if (this.cache.size >= this.config.cacheSize) {
      const firstKey = this.cache.keys().next().value;
      if (firstKey) {
        this.cache.delete(firstKey);
      }
    }
    this.cache.set(query, expanded);
  }

  /**
   * Clears all cached query expansions.
   */
  clearCache(): void {
    this.cache.clear();
    logInfo("QueryExpander: Cache cleared");
  }

  /**
   * Gets the current number of cached expansions.
   * @returns The size of the cache
   */
  getCacheSize(): number {
    return this.cache.size;
  }
}
