# Copyright (C) 2026 GoodData Corporation

"""Vocabulary harvesting utilities for rules retrieval.

Goal: improve lexical matching without maintaining a large manual alias list.
We harvest identifiers/paths/headings and derive normalized tokens that BM25 can match.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

# Minimal markdown punctuation stripping; keep underscores/hyphens for identifiers.
_MARKDOWN_PUNCT_RE = re.compile(r"[`#*_\[\]()]", re.MULTILINE)

# Inline code spans: `foo`
_INLINE_CODE_RE = re.compile(r"`([^`]+)`")

# Fenced code blocks: ```lang?\n...\n```
_CODE_FENCE_RE = re.compile(r"```[^\n]*\n([\s\S]*?)```", re.MULTILINE)

# Rough path-like tokens: contains at least one '/' and a file-ish suffix.
_PATH_RE = re.compile(r"\b[\w./-]+/[\w./-]+\.[a-zA-Z0-9]{1,6}\b")

# Dotted identifiers (packages, qualified names)
_DOTTED_IDENT_RE = re.compile(r"\b[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*){2,}\b")

# CamelCase identifiers (AntichainHandler, QtVisitor, etc.)
_CAMEL_RE = re.compile(r"\b[A-Z][a-z0-9]+(?:[A-Z][a-z0-9]+)+\b")

# snake_case / kebab-case-ish identifiers
_SNAKE_KEBAB_RE = re.compile(r"\b[a-zA-Z][a-zA-Z0-9_-]{2,}\b")


def _split_camel(token: str) -> list[str]:
    # Split CamelCase into parts, keeping digits as separate parts.
    parts = re.findall(r"[A-Z]+(?![a-z])|[A-Z]?[a-z]+|\d+", token)
    return [p for p in parts if p]


def _split_snake_kebab(token: str) -> list[str]:
    parts = re.split(r"[_-]+", token)
    return [p for p in parts if p]


def _normalize_token(token: str) -> str:
    return token.strip().lower()


def _emit_variants(raw: str) -> set[str]:
    """Emit normalized variants for an identifier/path-like token.

    Includes:
    - full normalized token
    - camel/snake/kebab parts
    """
    token = _normalize_token(raw)
    if not token:
        return set()

    out: set[str] = {token}

    # If token has path segments, include basename and segments.
    if "/" in token:
        segments = [s for s in token.split("/") if s]
        if segments:
            out.update(segments)
            out.add(segments[-1])
            # file name without extension
            if "." in segments[-1]:
                out.add(segments[-1].rsplit(".", 1)[0])

    # camel parts
    if re.search(r"[A-Z]", raw):
        parts = [_normalize_token(p) for p in _split_camel(raw)]
        out.update(parts)
        # join parts to match tokenize_query output for CamelCase in queries (e.g. QtVisitor -> qtvisitor)
        out.add("".join(parts))

    # snake/kebab parts
    if "_" in token or "-" in token:
        parts = [_normalize_token(p) for p in _split_snake_kebab(token)]
        out.update(parts)
        out.add("".join(parts))

    # Filter very short parts (keep full token even if short-ish).
    return {t for t in out if len(t) > 2} | ({token} if len(token) > 2 else set())


def extract_inline_code(text: str) -> list[str]:
    return [m.group(1) for m in _INLINE_CODE_RE.finditer(text)]


def extract_code_fences(text: str) -> list[str]:
    return [m.group(1) for m in _CODE_FENCE_RE.finditer(text)]


def extract_paths(text: str) -> list[str]:
    return [m.group(0) for m in _PATH_RE.finditer(text)]


def extract_identifier_like(text: str) -> list[str]:
    # Priority: dotted identifiers and camelcase, then general snake/kebab words.
    out: list[str] = []
    out.extend(m.group(0) for m in _DOTTED_IDENT_RE.finditer(text))
    out.extend(m.group(0) for m in _CAMEL_RE.finditer(text))
    out.extend(m.group(0) for m in _SNAKE_KEBAB_RE.finditer(text))
    return out


def tokenize_plain_text(text: str) -> set[str]:
    """Tokenize markdown-ish text into lowercase keyword tokens."""
    cleaned = _MARKDOWN_PUNCT_RE.sub(" ", text.lower())
    tokens = set(re.findall(r"[a-z0-9][\w-]*", cleaned))
    return {t for t in tokens if len(t) > 2}


@dataclass(frozen=True)
class HarvestedTokens:
    """Tokens for lexical scoring; boost_tokens get higher tf weight in BM25."""

    tokens: set[str]
    boost_tokens: set[str]


def harvest_section_tokens(*, heading_path: str | None, text: str) -> HarvestedTokens:
    """Harvest tokens from a markdown section.

    - tokens: baseline word tokens + identifier/path variants
    - boost_tokens: heading + identifier/path variants (strong retrieval anchors)
    """
    tokens = tokenize_plain_text(text)
    boost: set[str] = set()

    if heading_path:
        heading_tokens = tokenize_plain_text(heading_path)
        tokens |= heading_tokens
        boost |= heading_tokens

    # Inline + fenced code are strong anchors for dev queries
    code_blobs = extract_inline_code(text) + extract_code_fences(text)
    path_blobs = extract_paths(text)

    for blob in code_blobs + path_blobs:
        for ident in extract_identifier_like(blob):
            variants = _emit_variants(ident)
            tokens |= variants
            boost |= variants

    # Also run identifier extraction on the whole text for CamelCase mentions outside code
    for ident in extract_identifier_like(text):
        variants = _emit_variants(ident)
        tokens |= variants
        # don't boost everything from prose; only boost longer identifiers
        if len(ident) >= 6:
            boost |= variants

    return HarvestedTokens(tokens=tokens, boost_tokens=boost)
