modules/post_processor.py

import re
import logging
from modules.link_verifier import LinkVerifier
from modules.narrator_verifier import NarratorVerifier

REQUIRED_SECTIONS = [
    'تمهيد',
    'المبحث الأول',
    'المبحث الثاني',
    'المبحث الثالث',
    'المبحث الرابع',
    'المبحث الخامس',
    'قائمة المصادر',
]

SECTION_ALIASES = {
    'المقدمة': 'تمهيد',
    'التمهيد': 'تمهيد',
    'المبحث الاول': 'المبحث الأول',
    'المبحث الثاني': 'المبحث الثاني',
    'المبحث الثالث': 'المبحث الثالث',
    'المبحث الرابع': 'المبحث الرابع',
    'المبحث الخامس': 'المبحث الخامس',
    'الخاتمة': 'المبحث الخامس',
    'المصادر': 'قائمة المصادر',
    'المراجع': 'قائمة المصادر',
}


class PostProcessor:
    def __init__(self, term: str, category: str, draft: dict):
        self.term = term
        self.category = category
        self.draft = {"by_section": {}, **draft}
        self.verifier = LinkVerifier()
        self.verifier.build_index_from_draft(self.draft)
        self.valid_links = self._collect_valid_links()

    def process(self, text: str) -> str:
        text = self._clean_basic(text)
        text = self._remove_ai_headers(text)
        text = self._normalize_section_names(text)
        text = self._remove_duplicate_headers(text)
        text = self._remove_markdown_tables(text)
        # B4: ensure structure FIRST so placeholders exist before link fixing
        text = self._ensure_structure(text)
        text = self._fix_links(text)
        # B5: run narrator verification and append any disclaimer
        text = self._verify_narrators(text)
        return text

    def _remove_markdown_tables(self, text: str) -> str:
        """Convert any markdown tables into clean academic lists/paragraphs, banning table grids completely."""
        lines = text.split("\n")
        new_lines = []
        headers = []
        for line in lines:
            stripped = line.strip()
            if stripped.startswith("|") and stripped.endswith("|"):
                parts = [p.strip() for p in stripped.split("|")[1:-1]]
                # Ignore separator row like |---|---|
                if all(re.match(r"^:?-+:?$", p) for p in parts if p):
                    continue
                if not headers:
                    headers = parts
                    continue
                # Format table row into neat list item
                if len(parts) >= 2:
                    name = parts[1] if len(parts) > 1 else parts[0]
                    quote = parts[2] if len(parts) > 2 else ""
                    extra = f" ({parts[3]})" if len(parts) > 3 and parts[3] else ""
                    item_str = f"* **الراوي {name}**: {quote}{extra}".strip()
                    new_lines.append(item_str)
            else:
                headers = []
                new_lines.append(line)
        return "\n".join(new_lines)

    def _collect_valid_links(self) -> set[str]:
        links = set()
        for sec_data in self.draft.get("by_section", {}).values():
            for item in sec_data.get("results", []):
                link = item.get("link", "")
                if link:
                    links.add(link)
        return links

    # ── Basic cleanup ──────────────────────────────────────

    def _clean_basic(self, text: str) -> str:
        text = re.sub(r'^```(?:markdown)?\s*\n', '', text)
        text = re.sub(r'\n```\s*$', '', text)
        text = re.sub(r'\s*🔗\s*', ' 🔗 ', text)
        text = re.sub(r'\n{3,}', '\n\n', text)
        text = re.sub(r'[\u0000-\u0008\u000b\u000c\u000e-\u001f]', '', text)
        
        # Clean AI multilingual placeholders and English artifacts
        text = text.replace("某م", "راوٍ")
        text = text.replace("某某", "أحد")
        text = text.replace("某", "راوٍ")
        text = re.sub(r'\balongside\b', 'جنباً إلى جنب مع', text, flags=re.IGNORECASE)
        text = re.sub(r'\bcriterions\b', 'معايير', text, flags=re.IGNORECASE)
        text = re.sub(r'\bcriteria\b', 'معايير', text, flags=re.IGNORECASE)
        text = re.sub(r'\blinguistically grounded\b', 'مؤسساً لغوياً', text, flags=re.IGNORECASE)
        
        # Clean English placeholder leaks
        for pat in [r'without\s+citation', r'needs?\s+citation', r'requires?\s+citation', r'missing\s+citation', r'citation\s+needed']:
            text = re.sub(rf'\(\s*{pat}\s*\)', '(يحتاج توثيقاً)', text, flags=re.IGNORECASE)
        for pat in [r'needs?\s+publisher(?:\s+data|\s+info)?', r'needs?\s+publishing(?:\s+data|\s+info)?', r'missing\s+publisher(?:\s+data|\s+info)?', r'needs?\s+publication(?:\s+data|\s+info)?']:
            text = re.sub(rf'\(\s*{pat}\s*\)', '(يحتاج استكمال بيانات النشر)', text, flags=re.IGNORECASE)
        
        return text.strip()

    # ── Remove AI-generated section headers ─────────────────

    def _remove_ai_headers(self, text: str) -> str:
        """Remove headers that AI writes but are added programmatically."""
        lines = text.split('\n')
        result = []
        skip_next = False

        for i, line in enumerate(lines):
            stripped = line.strip()

            if skip_next:
                skip_next = False
                continue

            if stripped.startswith('## '):
                header_content = stripped[3:].strip()
                next_idx = i + 1
                while next_idx < len(lines) and lines[next_idx].strip() == '':
                    next_idx += 1

                if next_idx < len(lines):
                    next_line = lines[next_idx].strip()
                    if len(next_line) < 100 and (
                        next_line in header_content or
                        header_content in next_line or
                        any(alias in next_line for alias in SECTION_ALIASES.keys())
                    ):
                        continue

            result.append(lines[i])

        return '\n'.join(result)

    # ── Section name normalization ──────────────────────

    def _normalize_section_names(self, text: str) -> str:
        text = self._apply_alias(text, 'التمهيد', '## تمهيد')
        text = self._apply_alias(text, 'تمهيد', '## تمهيد')
        text = self._apply_alias(text, 'المقدمة', '## تمهيد')
        for i in range(1, 6):
            arabic = self._to_arabic(i)
            variations = [arabic]
            if "أ" in arabic:
                variations.append(arabic.replace("أ", "ا"))
            if "إ" in arabic:
                variations.append(arabic.replace("إ", "ا"))
                
            for var in variations:
                text = self._apply_alias(
                    text,
                    f'المبحث {var}',
                    f'## المبحث {arabic}',
                )
            text = self._apply_alias(
                text,
                f'المبحث {i}',
                f'## المبحث {arabic}',
            )
            for var in variations:
                text = re.sub(
                    rf'(?:^|\n)\s*المبحث\s*{var}\s*[:\-–]',
                    f'\n## المبحث {arabic}:',
                    text,
                )
        text = self._apply_alias(text, 'قائمة المصادر والمراجع', '## قائمة المصادر')
        text = self._apply_alias(text, 'قائمة المصادر', '## قائمة المصادر')
        text = self._apply_alias(text, 'المصادر والمراجع', '## قائمة المصادر')
        text = self._apply_alias(text, 'المصادر', '## قائمة المصادر')
        text = self._apply_alias(text, 'المراجع', '## قائمة المصادر')
        return text

    def _apply_alias(self, text: str, alias: str, replacement: str) -> str:
        pattern = re.compile(
            r'(?:^|\n)\s*(?:#+\s*|\*{0,2}\s*)'
            + re.escape(alias)
            + r'\s*\*{0,2}\s*(?:[:\-–]+\s*)?',
            re.MULTILINE,
        )
        return pattern.sub(f'\n{replacement}\n', text)

    # ── Remove duplicate headers ─────────────────────────

    def _remove_duplicate_headers(self, text: str) -> str:
        """Remove consecutive duplicate ## headers."""
        lines = text.split('\n')
        result = []
        prev_header = ""

        for line in lines:
            stripped = line.strip()
            if stripped.startswith('## '):
                if stripped == prev_header:
                    continue
                prev_header = stripped
            elif stripped:
                prev_header = ""
            result.append(line)

        return '\n'.join(result)

    # ── Link fixing ─────────────────────────────────────

    def _fix_links(self, text: str) -> str:
        # Call the live self-healing verifier to fix or placeholderize links
        text = self.verifier.verify_and_fix(text, self.valid_links)
        return text

    # ── Structure enforcement ───────────────────────────

    def _ensure_structure(self, text: str) -> str:
        existing = set()
        for line in text.split('\n'):
            stripped = line.strip()
            if stripped.startswith('#') or stripped.startswith('**'):
                for sec in REQUIRED_SECTIONS:
                    if sec in stripped:
                        existing.add(sec)

        missing = [s for s in REQUIRED_SECTIONS if s not in existing]

        if missing:
            logging.warning(
                f"⚠️ الدراسة «{self.term}» تفتقد: {missing}"
            )
            for sec in missing:
                text += f'\n\n## {sec}\n[لم يُكتب — يحتاج إضافة يدوية]'

        text = re.sub(r'قائمة\s+قائمة\s+المصادر', 'قائمة المصادر', text)
        text = re.sub(r'قائمة المصادر\s+وقائمة\s+المصادر', 'قائمة المصادر', text)

        return text

    # ── Narrator Verification ───────────────────────────────────────────

    def _verify_narrators(self, text: str) -> str:
        """Run project-wide narrator verification and append disclaimer if needed."""
        try:
            verifier = NarratorVerifier(self.term, self.category)
            result = verifier.verify(text)
            if result.has_issues:
                disclaimer = verifier.build_disclaimer(result)
                if disclaimer:
                    logging.warning(
                        f"⚠️ [{self.term}] Narrator attribution issues: "
                        f"{len(result.ambiguous)} ambiguous, {len(result.rejected)} rejected"
                    )
                    # Append warning note before قائمة المصادر section
                    insertion = f"\n\n> **ملاحظة منهجية:** {disclaimer}\n"
                    # Insert before sources list
                    masadir_match = re.search(r"\n## قائمة المصادر", text)
                    if masadir_match:
                        pos = masadir_match.start()
                        text = text[:pos] + insertion + text[pos:]
                    else:
                        text += insertion
        except Exception as e:
            logging.debug(f"NarratorVerifier skipped for {self.term}: {e}")
        return text

    # ── Re-save study after processing ──────────────────

    def re_save(self, text: str, file_path) -> None:
        file_path.write_text(text, encoding="utf-8")

    def _to_arabic(self, n: int) -> str:
        arabic = ['','الأول','الثاني','الثالث','الرابع','الخامس']
        return arabic[n] if n < len(arabic) else str(n)