Blame for backend/python-service/comment_target_resolver.py

from __future__ import annotations

import re
import unicodedata
from dataclasses import dataclass
from typing import Optional

from docx_text import extract_docx_paragraphs
from models import ParagraphText, ReviewFinding, ResolveStrategy
from pathlib import Path


@dataclass
class CharMap:
    normalized: str
    orig_offsets: list[int]


@dataclass
class ResolveResult:
    target_text: str
    before_anchor: str
    after_anchor: str
    strategy: ResolveStrategy
    confidence: float
    paragraph_id: str
    was_resolved: bool = False
    message: str = ""


def _build_char_map(text: str) -> CharMap:
    norm_chars: list[str] = []
    offsets: list[int] = []
    for i, ch in enumerate(text):
        folded = _fold_char(ch)
        if folded is not None:
            norm_chars.append(folded)
            offsets.append(i)
    return CharMap(normalized="".join(norm_chars), orig_offsets=offsets)


def _fold_char(ch: str) -> Optional[str]:
    if ch == "\u00a0":
        return None
    cat = unicodedata.category(ch)
    if cat.startswith("Z"):
        return None
    if ch in "\u2018\u2019\u201b":
        return "'"
    if ch in "\u201c\u201d\u201e":
        return '"'
    if ch in "\u2013\u2014\u2015\u2212\u2010":
        return "-"
    if ch == "\u00ad":
        return None
    return ch


def _normalize_for_match(text: str) -> str:
    return _build_char_map(text).normalized


def _find_all(haystack: str, needle: str) -> list[int]:
    positions: list[int] = []
    start = 0
    while True:
        idx = haystack.find(needle, start)
        if idx == -1:
            break
        positions.append(idx)
        start = idx + 1
    return positions


def _map_norm_range_to_orig(char_map: CharMap, norm_start: int, norm_end: int) -> tuple[int, int]:
    orig_start = char_map.orig_offsets[norm_start]
    if norm_end >= len(char_map.orig_offsets):
        orig_end = char_map.orig_offsets[-1] + 1
    else:
        orig_end = char_map.orig_offsets[norm_end - 1] + 1
    return orig_start, orig_end


def _char_bigrams(text: str) -> set[str]:
    normalized = re.sub(r"\s+", "", text.lower())
    if len(normalized) < 2:
        return set()
    return {normalized[i:i+2] for i in range(len(normalized) - 1)}


def _bigram_containment(needle: set[str], haystack: set[str]) -> float:
    if not needle:
        return 0.0
    return len(needle & haystack) / len(needle)

def _sentence_tokenize(text: str) -> list[tuple[int, int]]:
    pattern = re.compile(r"(?<=[.!?])\s+")
    sentences: list[tuple[int, int]] = []
    last = 0
    for m in pattern.finditer(text):
        end = m.start()
        if end > last:
            sentences.append((last, end))
        last = m.end()
    if last < len(text):
        sentences.append((last, len(text)))
    return sentences

def generate_editable_subspan_candidates(
    para_text: str,
    target_text: str,
) -> list[str]:
    """
    Generates smaller text windows from the target text or surrounding paragraph
    to find an editable subspan when the full target crosses complex XML containers.
    """
    candidates = []
    
    # We must find where the target text is in the paragraph
    idx = para_text.find(target_text)
    if idx == -1:
        # If not exact match, just use the target text itself for fallback logic
        working_text = target_text
        center = len(target_text) // 2
    else:
        working_text = para_text
        center = idx + (len(target_text) // 2)

    # 1. Try full sentence
    sentences = _sentence_tokenize(working_text)
    for s_start, s_end in sentences:
        if s_start <= center <= s_end:
            sent = working_text[s_start:s_end].strip()
            if len(sent.split()) >= 5:
                candidates.append(sent)

    # 2. Try clause windows
    if idx != -1:
        clause_text = working_text[idx:idx+len(target_text)]
    else:
        clause_text = working_text
        
    clauses = re.split(r"[,;:\(\)\-]", clause_text)
    # Sort clauses by length, prefer ones with 8-20 words
    scored_clauses = []
    for c in clauses:
        c = c.strip()
        words = c.split()
        if len(words) < 5:
            continue
        # We want to prefer clauses that overlap the center, but within the target it's all good.
        score = 0
        if 8 <= len(words) <= 20:
            score += 10
        scored_clauses.append((score, c))
        
    scored_clauses.sort(key=lambda x: x[0], reverse=True)
    for _, c in scored_clauses:
        candidates.append(c)
        
    # 3. Sliding word windows
    words = clause_text.split()
    if len(words) > 8:
        for window_size in [20, 16, 12, 8]:
            if window_size >= len(words):
                continue
            
            # Start from center-ish
            mid_word = len(words) // 2
            
            # Generate windows shifting outwards from center
            # Simple approach: just generate all windows of this size and sort by distance to center
            windows = []
            for i in range(len(words) - window_size + 1):
                w_center = i + (window_size // 2)
                dist = abs(w_center - mid_word)
                windows.append((dist, " ".join(words[i:i+window_size])))
                
            windows.sort(key=lambda x: x[0])
            for _, w in windows:
                candidates.append(w)
                
    # Deduplicate while preserving order
    seen = set()
    unique_candidates = []
    for c in candidates:
        if c not in seen:
            seen.add(c)
            # Only use candidates that appear exactly once in the paragraph
            if para_text.count(c) == 1:
                unique_candidates.append(c)
                
    return unique_candidates

def _search_all_paragraphs(
    paragraphs: list[ParagraphText],
    target_text: str,
    before_anchor: str,
    after_anchor: str,
) -> Optional[tuple[ResolveResult, str]]:
    best: Optional[tuple[ResolveResult, str]] = None
    for para in paragraphs:
        result = _resolve_target_in_paragraph(
            para_text=para.text,
            target_text=target_text,
            before_anchor=before_anchor,
            after_anchor=after_anchor,
        )
        if result is None:
            continue
        if result.strategy == "exact" and result.confidence == 1.0:
            return (result, para.id)
        if best is None or result.confidence > best[0].confidence:
            best = (result, para.id)
    return best


@dataclass
class FindingResolution:
    finding: ReviewFinding
    original_target_text: str
    original_before_anchor: str
    original_after_anchor: str
    strategy: ResolveStrategy
    confidence: float
    was_resolved: bool


def resolve_findings_against_docx(
    input_docx: Path,
    findings: list[ReviewFinding],
    allow_paragraph_fallback: bool = False,
) -> list[FindingResolution]:
    paragraphs = extract_docx_paragraphs(input_docx)
    para_map: dict[str, ParagraphText] = {p.id: p for p in paragraphs}

    resolved: list[FindingResolution] = []

    for finding in findings:
        para = para_map.get(finding.start_paragraph_id)

        # Try the specified paragraph first
        if para is not None:
            result = _resolve_target_in_paragraph(
                para_text=para.text,
                target_text=finding.target_text,
                before_anchor=finding.before_anchor,
                after_anchor=finding.after_anchor,
            )
            if result is not None and (result.strategy == "exact" or result.was_resolved):
                if result.was_resolved:
                    new_finding = finding.model_copy(update={
                        "target_text": result.target_text,
                        "before_anchor": result.before_anchor,
                        "after_anchor": result.after_anchor,
                    })
                    resolved.append(FindingResolution(
                        finding=new_finding,
                        original_target_text=finding.target_text,
                        original_before_anchor=finding.before_anchor,
                        original_after_anchor=finding.after_anchor,
                        strategy=result.strategy,
                        confidence=result.confidence,
                        was_resolved=True,
                    ))
                else:
                    resolved.append(FindingResolution(
                        finding=finding,
                        original_target_text=finding.target_text,
                        original_before_anchor=finding.before_anchor,
                        original_after_anchor=finding.after_anchor,
                        strategy="exact",
                        confidence=1.0,
                        was_resolved=False,
                    ))
                continue

        # Search across all paragraphs when the specified paragraph
        # didn't match or had wrong ID (common: paragraph ID mismatch
        # between what the LLM sees and what the Python extractor produces)
        global_result = _search_all_paragraphs(
            paragraphs,
            finding.target_text,
            finding.before_anchor,
            finding.after_anchor,
        )
        if global_result is not None:
            result, matched_para_id = global_result
            if result.was_resolved:
                new_finding = finding.model_copy(update={
                    "target_text": result.target_text,
                    "before_anchor": result.before_anchor,
                    "after_anchor": result.after_anchor,
                    "start_paragraph_id": matched_para_id,
                    "end_paragraph_id": matched_para_id,
                })
                resolved.append(FindingResolution(
                    finding=new_finding,
                    original_target_text=finding.target_text,
                    original_before_anchor=finding.before_anchor,
                    original_after_anchor=finding.after_anchor,
                    strategy=result.strategy,
                    confidence=result.confidence,
                    was_resolved=True,
                ))
            else:
                new_finding = finding.model_copy(update={
                    "start_paragraph_id": matched_para_id,
                    "end_paragraph_id": matched_para_id,
                })
                resolved.append(FindingResolution(
                    finding=new_finding,
                    original_target_text=finding.target_text,
                    original_before_anchor=finding.before_anchor,
                    original_after_anchor=finding.after_anchor,
                    strategy="exact",
                    confidence=1.0,
                    was_resolved=False,
                ))
            continue
        elif allow_paragraph_fallback and para is not None and para.text:
            first_sentence_end = para.text.find(".")
            if first_sentence_end == -1:
                first_sentence_end = min(len(para.text), 200)
            else:
                first_sentence_end += 1
            anchor_text = para.text[:first_sentence_end]
            new_finding = finding.model_copy(update={
                "target_text": anchor_text,
                "before_anchor": "",
                "after_anchor": "",
            })
            resolved.append(FindingResolution(
                finding=new_finding,
                original_target_text=finding.target_text,
                original_before_anchor=finding.before_anchor,
                original_after_anchor=finding.after_anchor,
                strategy="paragraph_fallback",
                confidence=0.4,
                was_resolved=True,
            ))
        else:
            resolved.append(FindingResolution(
                finding=finding,
                original_target_text=finding.target_text,
                original_before_anchor=finding.before_anchor,
                original_after_anchor=finding.after_anchor,
                strategy="not_resolved",
                confidence=0.0,
                was_resolved=False,
            ))

    return resolved


def resolve_single_target(
    paragraphs: list[ParagraphText],
    block_id: str,
    target_text: str,
    before_anchor: str,
    after_anchor: str,
    allow_paragraph_fallback: bool = False,
) -> ResolveResult:
    para = next((p for p in paragraphs if p.id == block_id), None)
    if para is None:
        return ResolveResult(
            target_text=target_text,
            before_anchor=before_anchor,
            after_anchor=after_anchor,
            strategy="not_resolved",
            confidence=0.0,
            paragraph_id=block_id,
            message=f"Paragraph {block_id} not found in document.",
        )

    result = _resolve_target_in_paragraph(
        para_text=para.text,
        target_text=target_text,
        before_anchor=before_anchor,
        after_anchor=after_anchor,
    )

    if result is not None:
        result.paragraph_id = block_id
        return result

    if allow_paragraph_fallback and para.text:
        first_dot = para.text.find(".")
        end = first_dot + 1 if first_dot != -1 else min(len(para.text), 200)
        return ResolveResult(
            target_text=para.text[:end],
            before_anchor="",
            after_anchor="",
            strategy="paragraph_fallback",
            confidence=0.4,
            paragraph_id=block_id,
            was_resolved=True,
        )

    return ResolveResult(
        target_text=target_text,
        before_anchor=before_anchor,
        after_anchor=after_anchor,
        strategy="not_resolved",
        confidence=0.0,
        paragraph_id=block_id,
        message="Could not resolve target text against paragraph.",
    )
def _resolve_target_in_paragraph(
    para_text: str,
    target_text: str,
    before_anchor: str,
    after_anchor: str,
) -> Optional[ResolveResult]:
    para_map = _build_char_map(para_text)

    # --- Layer 1: Exact match ---
    exact_positions = _find_all(para_text, target_text)
    if len(exact_positions) == 1:
        return ResolveResult(
            target_text=target_text,
            before_anchor=before_anchor,
            after_anchor=after_anchor,
            strategy="exact",
            confidence=1.0,
            paragraph_id="",
        )
    if len(exact_positions) > 1:
        if before_anchor or after_anchor:
            for pos in exact_positions:
                if before_anchor and para_text[:pos].endswith(before_anchor):
                    return ResolveResult(
                        target_text=target_text,
                        before_anchor=before_anchor,
                        after_anchor=after_anchor,
                        strategy="exact",
                        confidence=1.0,
                        paragraph_id="",
                    )
                if after_anchor and para_text[pos + len(target_text):].startswith(after_anchor):
                    return ResolveResult(
                        target_text=target_text,
                        before_anchor=before_anchor,
                        after_anchor=after_anchor,
                        strategy="exact",
                        confidence=1.0,
                        paragraph_id="",
                    )

    # --- Layer 2: Whitespace-insensitive match ---
    norm_target = _normalize_for_match(target_text)
    if norm_target:
        norm_positions = _find_all(para_map.normalized, norm_target)

        # If no exact case match, try case-insensitive
        if not norm_positions:
            norm_lower = norm_target.lower()
            para_lower = para_map.normalized.lower()
            norm_positions_ci = _find_all(para_lower, norm_lower)
            if norm_positions_ci and len(norm_positions_ci) == 1:
                norm_positions = norm_positions_ci

        if len(norm_positions) == 1:
            ns = norm_positions[0]
            ne = ns + len(norm_target)
            orig_start, orig_end = _map_norm_range_to_orig(para_map, ns, ne)
            resolved = para_text[orig_start:orig_end]
            return ResolveResult(
                target_text=resolved,
                before_anchor="",
                after_anchor="",
                strategy="whitespace_normalized",
                confidence=0.95,
                paragraph_id="",
                was_resolved=True,
            )

        if len(norm_positions) > 1:
            # --- Layer 3: Anchor-aware normalized match ---
            norm_before = _normalize_for_match(before_anchor) if before_anchor else ""
            norm_after = _normalize_for_match(after_anchor) if after_anchor else ""

            best_idx: Optional[int] = None
            best_score = -1

            for pos in norm_positions:
                end_pos = pos + len(norm_target)
                score = 0

                if norm_before:
                    prefix = para_map.normalized[max(0, pos - len(norm_before)):pos]
                    if prefix.lower().endswith(norm_before.lower()):
                        score += 2
                    elif prefix[-len(norm_before):].lower() == norm_before[-len(prefix):].lower() if len(prefix) >= len(norm_before) else prefix.lower() == norm_before[-len(prefix):].lower():
                        score += 1

                if norm_after:
                    suffix_start = end_pos
                    suffix = para_map.normalized[suffix_start:suffix_start + len(norm_after)]
                    if suffix.lower().startswith(norm_after.lower()):
                        score += 2
                    elif suffix[:len(norm_after)].lower() == norm_after[:len(suffix)].lower() if len(suffix) >= len(norm_after) else suffix.lower() == norm_after[:len(suffix)].lower():
                        score += 1

                if score > best_score:
                    best_score = score
                    best_idx = pos

            if best_idx is not None and best_score > 0:
                ns = best_idx
                ne = ns + len(norm_target)
                orig_start, orig_end = _map_norm_range_to_orig(para_map, ns, ne)
                resolved = para_text[orig_start:orig_end]
                return ResolveResult(
                    target_text=resolved,
                    before_anchor="",
                    after_anchor="",
                    strategy="anchor_normalized",
                    confidence=0.8 if best_score >= 3 else 0.65,
                    paragraph_id="",
                    was_resolved=True,
                )

    # --- Layer 4: Sentence-level fallback ---
    sentences = _sentence_tokenize(para_text)
    target_bigrams = _char_bigrams(target_text)
    if target_bigrams and sentences:
        scored: list[tuple[float, int, int]] = []
        for s_start, s_end in sentences:
            sent = para_text[s_start:s_end]
            sent_bigrams = _char_bigrams(sent)
            score = _bigram_containment(target_bigrams, sent_bigrams)
            scored.append((score, s_start, s_end))
        scored.sort(key=lambda x: x[0], reverse=True)
        if scored:
            best = scored[0]
            runner_up = scored[1][0] if len(scored) > 1 else 0.0
            if best[0] >= 0.70 and best[0] - runner_up > 0.10:
                resolved = para_text[best[1]:best[2]]
                return ResolveResult(
                    target_text=resolved,
                    before_anchor="",
                    after_anchor="",
                    strategy="sentence_fallback",
                    confidence=round(best[0], 2),
                    paragraph_id="",
                    was_resolved=True,
                )

    return None