"""PDF text extraction and section detection with ecology-aware header matching.

Uses pymupdf4llm for high-quality extraction (handles multi-column layouts,
tables, and complex formatting) with fallback to raw PyMuPDF.
"""

from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Optional

import fitz  # PyMuPDF

from res_sum.schemas import DomainConfig

logger = logging.getLogger("res_sum")


# ---------------------------------------------------------------------------
# Default section headers (always recognised)
# ---------------------------------------------------------------------------

_BASE_HEADERS: dict[str, list[str]] = {
    "introduction": ["introduction"],
    "methodology": [
        "methodology",
        "methods",
        "materials and methods",
        "materials & methods",
        "material and methods",
    ],
    "results": ["results", "results and discussion"],
    "discussion": ["discussion"],
    "conclusion": ["conclusion", "conclusions", "concluding remarks"],
}

_STOP_HEADERS = [
    "references",
    "acknowledgements",
    "acknowledgments",
    "acknowledgement",
    "acknowledgment",
    "supplementary",
    "supplemental",
    "supplementary material",
    "supplementary materials",
    "appendix",
    "appendices",
    "supporting information",
    "data availability",
    "data accessibility",
    "author contributions",
    "conflict of interest",
    "conflicts of interest",
    "competing interests",
    "funding",
]


def _build_header_pattern(headers: list[str]) -> re.Pattern:
    """Build a regex that matches section headers with optional numbering.

    Matches patterns like:
      - "Introduction"
      - "2. Methods"
      - "2.1 Materials and Methods"
      - "II. RESULTS"
      - "METHODS" (all caps)
      - "# Introduction" (Markdown from pymupdf4llm)
      - "## 2.1 Materials and Methods" (Markdown with numbering)
      - "**Methods**" (bold Markdown)
    """
    escaped = [re.escape(h) for h in headers]
    joined = "|".join(escaped)
    # Optional: markdown heading prefix (# or ##), bold markers (**),
    # numbering (1. or 2.1 or II.), then header text.
    # Allow trailing punctuation or inline content (e.g., "ACKNOWLEDGMENTS. J. J. Kolbe...")
    pattern = (
        r"^\s*"
        r"(?:#{1,4}\s*)?"              # Optional Markdown heading (# to ####)
        r"(?:\*{1,2})?"                # Optional bold start (* or **)
        r"(?:\d+\.?\d*\.?\s*|[IVXLC]+\.?\s*)?"  # Optional numbering
        rf"({joined})"                 # The header text
        r"(?:\*{1,2})?"                # Optional bold end
        r"(?:[.:\s].*)?$"             # Optional trailing punctuation/content
    )
    return re.compile(pattern, re.IGNORECASE)


def extract_text_from_pdf(pdf_path: str | Path) -> str:
    """Extract text from a PDF file.

    Uses pymupdf4llm for high-quality Markdown extraction (handles multi-column
    layouts, tables, and complex formatting). Falls back to raw PyMuPDF if
    pymupdf4llm encounters an error.

    Args:
        pdf_path: Path to the PDF file.

    Returns:
        Extracted text as a single string.
    """
    pdf_path = Path(pdf_path)
    if not pdf_path.exists():
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    # Try pymupdf4llm first (better multi-column, table handling)
    try:
        import pymupdf4llm
        text = pymupdf4llm.to_markdown(str(pdf_path))
        if text and text.strip():
            return text
    except Exception as e:
        logger.debug(f"pymupdf4llm failed for '{pdf_path.name}': {e}. Falling back to PyMuPDF.")

    # Fallback: raw PyMuPDF extraction
    try:
        doc = fitz.open(str(pdf_path))
        pages = [page.get_text() for page in doc]
        doc.close()
        return "\n".join(pages)
    except Exception as e:
        raise RuntimeError(f"Error reading PDF '{pdf_path.name}': {e}") from e


def extract_sections(
    text: str,
    start_section: str = "methodology",
    domain_config: Optional[DomainConfig] = None,
) -> tuple[str, dict[str, str]]:
    """Extract and combine sections from a research paper's text.

    Args:
        text: Full text of the paper.
        start_section: Which section to start extracting from
            ("introduction" or "methodology").
        domain_config: Optional domain config with additional section headers.

    Returns:
        Tuple of (combined_text, sections_dict).
    """
    # Merge base headers with any domain-specific ones
    all_method_headers = list(_BASE_HEADERS["methodology"])
    if domain_config and domain_config.section_headers:
        all_method_headers.extend(domain_config.section_headers)

    # Build lookup: header text → canonical section name
    header_to_section: dict[str, str] = {}
    for section_name, headers in _BASE_HEADERS.items():
        for h in headers:
            header_to_section[h] = section_name
    # Domain-specific headers map to "methodology" (they're sub-sections of methods)
    if domain_config:
        for h in domain_config.section_headers:
            if h not in header_to_section:
                header_to_section[h] = "methodology"

    # Build patterns
    all_headers = list(header_to_section.keys())
    header_pattern = _build_header_pattern(all_headers)
    stop_pattern = _build_header_pattern(_STOP_HEADERS)

    sections: dict[str, str] = {
        "introduction": "",
        "methodology": "",
        "results": "",
        "discussion": "",
        "conclusion": "",
    }

    current_section: Optional[str] = None
    start_extracting = False
    found_sections: list[str] = []

    for line in text.split("\n"):
        line_stripped = line.strip()
        line_lower = line_stripped.lower()

        # Check for stop headers
        if stop_pattern.match(line_stripped):
            start_extracting = False
            current_section = None
            continue

        # Check for section headers
        header_match = header_pattern.match(line_stripped)
        if header_match:
            matched_text = header_match.group(1).lower()
            # Find which canonical section this maps to
            canonical = header_to_section.get(matched_text)
            if canonical is None:
                # Fuzzy match against known headers
                for known_header, section_name in header_to_section.items():
                    if known_header in matched_text or matched_text in known_header:
                        canonical = section_name
                        break

            if canonical:
                current_section = canonical
                if canonical not in found_sections:
                    found_sections.append(canonical)

                # Decide whether to start extracting based on start_section
                if start_section == "introduction":
                    start_extracting = True
                elif start_section == "methodology" and canonical != "introduction":
                    start_extracting = True
                elif canonical == start_section:
                    start_extracting = True
                continue

        # Accumulate text into current section
        if start_extracting and current_section:
            sections[current_section] += line + "\n"

    # Combine sections in order
    if start_section == "introduction":
        ordered = ["introduction", "methodology", "results", "discussion", "conclusion"]
    else:
        ordered = ["methodology", "results", "discussion", "conclusion"]

    combined_text = "".join(sections[s] for s in ordered)

    return combined_text, sections
