"""Citation post-processing for deep research reports.

Scans synthesized reports for inline [N] citations, verifies consistency
against the source registry, removes dangling references, and appends
a deterministic Sources section built from state rather than LLM output.
"""

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from foundry_mcp.core.research.models.deep_research import DeepResearchState

logger = logging.getLogger(__name__)

# Matches [N] where N is one or more digits, but NOT inside markdown link
# syntax like [text](url). The negative lookahead (?!\() ensures we skip
# patterns followed by a parenthesised URL.
_CITATION_RE = re.compile(r"\[(\d+)\](?!\()")


def extract_cited_numbers(report: str) -> set[int]:
    """Extract all citation numbers referenced in the report.

    Finds all ``[N]`` patterns in the report text. Ignores markdown
    link syntax by only matching bare numeric brackets.

    Args:
        report: The markdown report text.

    Returns:
        Set of cited integer citation numbers.
    """
    return {int(m.group(1)) for m in _CITATION_RE.finditer(report)}


def build_sources_section(
    state: "DeepResearchState",
    *,
    cited_only: bool = False,
    cited_numbers: set[int] | None = None,
) -> str:
    """Build a deterministic ``## Sources`` section from state.

    Args:
        state: Research state containing all sources with citation numbers.
        cited_only: If True, only include sources that were actually cited.
        cited_numbers: Pre-computed set of cited numbers (avoids re-scanning).

    Returns:
        Markdown string for the Sources section (including the heading).
    """
    citation_map = state.get_citation_map()
    if not citation_map:
        return ""

    lines = ["", "## Sources", ""]
    for cn in sorted(citation_map):
        if cited_only and (cited_numbers is None or cn not in cited_numbers):
            continue
        source = citation_map[cn]
        title = source.title or "Untitled"
        if source.url:
            lines.append(f"[{cn}] [{title}]({source.url})")
        else:
            lines.append(f"[{cn}] {title}")
    lines.append("")
    return "\n".join(lines)


def remove_dangling_citations(report: str, valid_numbers: set[int]) -> str:
    """Remove citation markers that reference non-existent sources.

    Replaces ``[N]`` with empty string when N is not in *valid_numbers*.

    Args:
        report: The markdown report text.
        valid_numbers: Set of citation numbers that exist in state.

    Returns:
        Report with dangling citations removed.
    """

    def _replace(match: re.Match) -> str:
        num = int(match.group(1))
        if num in valid_numbers:
            return match.group(0)
        return ""

    return _CITATION_RE.sub(_replace, report)


def strip_llm_sources_section(report: str) -> str:
    """Remove any Sources/References section generated by the LLM.

    Looks for common heading patterns (``## Sources``, ``## References``,
    ``## Works Cited``) and removes everything from that heading to the
    next heading of equal or higher level, or the end of the report.

    Args:
        report: The markdown report text.

    Returns:
        Report with LLM-generated sources section removed.
    """
    # Match ## Sources, ## References, ## Works Cited (case-insensitive)
    pattern = re.compile(
        r"^(#{1,2}\s+(?:Sources|References|Works\s+Cited|Bibliography))\s*$",
        re.MULTILINE | re.IGNORECASE,
    )
    match = pattern.search(report)
    if not match:
        return report

    start = match.start()
    heading_level = match.group(1).count("#")

    # Find the next heading of equal or higher level
    rest = report[match.end() :]
    next_heading = re.search(
        rf"^#{{{1},{heading_level}}}\s+\S",
        rest,
        re.MULTILINE,
    )
    if next_heading:
        end = match.end() + next_heading.start()
    else:
        end = len(report)

    # Strip and clean up trailing whitespace
    return report[:start].rstrip() + report[end:]


def postprocess_citations(
    report: str,
    state: "DeepResearchState",
) -> tuple[str, dict]:
    """Run full citation post-processing pipeline on a report.

    Steps:
    1. Extract all cited ``[N]`` numbers from the report.
    2. Remove any LLM-generated Sources section.
    3. Remove dangling citations (referencing non-existent sources).
    4. Append a deterministic Sources section from state.

    Args:
        report: The synthesized markdown report.
        state: Research state with all sources and citation numbers.

    Returns:
        Tuple of (processed_report, metadata_dict) where metadata contains
        citation statistics for audit logging.
    """
    citation_map = state.get_citation_map()
    valid_numbers = set(citation_map.keys())

    # 1. Extract cited numbers
    cited_numbers = extract_cited_numbers(report)

    # 2. Strip any LLM-generated sources section
    report = strip_llm_sources_section(report)

    # 3. Remove dangling citations
    dangling = cited_numbers - valid_numbers
    if dangling:
        logger.warning(
            "Removing %d dangling citation(s): %s",
            len(dangling),
            sorted(dangling),
        )
        report = remove_dangling_citations(report, valid_numbers)
        # Recompute after removal
        cited_numbers = extract_cited_numbers(report)

    # 4. Append deterministic Sources section
    sources_section = build_sources_section(state)
    if sources_section:
        report = report.rstrip() + "\n" + sources_section

    # Compute unreferenced sources (have citation number but never cited)
    unreferenced = valid_numbers - cited_numbers
    if unreferenced:
        logger.info(
            "%d source(s) have citation numbers but were not referenced in the report: %s",
            len(unreferenced),
            sorted(unreferenced),
        )

    metadata = {
        "total_citations_in_report": len(cited_numbers),
        "total_sources_with_numbers": len(valid_numbers),
        "dangling_citations_removed": len(dangling),
        "unreferenced_sources": len(unreferenced),
    }

    return report, metadata
