"""
lambdagent.ingest — citation-aware document ingestion (M2 Phase E).

CR-20260607-001 rev.2 §6.2 (FR-004/FR-005) + Q7 (Obsidian import).

The research agent packs (Phase F) promise "回答和报告中尽量保留来源文件、
页码或段落" (FR-005). To deliver that, ingestion must preserve provenance:
every chunk carries which file + which page it came from, so retrieval can
render `[Source N | paper.pdf p.12]`.

Three building blocks, all pure + dependency-light:
  - extract_pdf_pages(path)      → [(page_number, text), ...]   (1-indexed)
  - chunk_pages(pages, source)   → [CitedChunk, ...]            (carries page)
  - import_obsidian_vault(dir)   → [CitedChunk, ...]            (markdown vault)

PDF extraction prefers pdfplumber (better layout + reliable page split),
falls back to PyPDF2, and raises a clear error if neither is installed.
Obsidian/markdown needs no dependency.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple


@dataclass
class CitedChunk:
    """A retrievable chunk that remembers where it came from."""
    text: str
    source: str                 # filename or relative path
    page: Optional[int] = None  # 1-indexed PDF page, or None for non-paged
    chunk_index: int = 0        # position within the source
    extra: Dict[str, Any] = field(default_factory=dict)

    def to_metadata(self) -> Dict[str, Any]:
        """Metadata dict suitable for rag.SimpleVectorStore.add(metadata=...).
        The keys `source` + `page` are what rag's citation formatter looks
        for to render `[Source N | source p.page]`."""
        md: Dict[str, Any] = {"source": self.source, "chunk_index": self.chunk_index}
        if self.page is not None:
            md["page"] = self.page
        md.update(self.extra)
        return md


# ── PDF page extraction ───────────────────────────────────────────────────

def extract_pdf_pages(path: str) -> List[Tuple[int, str]]:
    """Extract text per page. Returns [(page_number_1indexed, text), ...].

    Prefers pdfplumber (cleaner layout-aware extraction), falls back to
    PyPDF2. Pages with no extractable text are still returned (empty string)
    so page numbering stays aligned with the physical document.
    """
    if not os.path.isfile(path):
        raise FileNotFoundError(f"PDF not found: {path}")

    # pdfplumber first
    try:
        import pdfplumber  # type: ignore
        out: List[Tuple[int, str]] = []
        with pdfplumber.open(path) as pdf:
            for i, page in enumerate(pdf.pages):
                out.append((i + 1, page.extract_text() or ""))
        return out
    except ImportError:
        pass

    # PyPDF2 fallback
    try:
        import PyPDF2  # type: ignore
        out = []
        with open(path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            for i, page in enumerate(reader.pages):
                out.append((i + 1, page.extract_text() or ""))
        return out
    except ImportError:
        raise RuntimeError(
            "PDF extraction needs pdfplumber or PyPDF2. "
            "Install with: pip install pdfplumber"
        )


# ── Page-preserving chunker ────────────────────────────────────────────────

def chunk_text(text: str, chunk_size: int = 800, overlap: int = 100) -> List[str]:
    """Split text into overlapping windows on paragraph/whitespace
    boundaries. Pure helper used by chunk_pages."""
    text = text.strip()
    if not text:
        return []
    if len(text) <= chunk_size:
        return [text]

    chunks: List[str] = []
    start = 0
    n = len(text)
    while start < n:
        end = min(start + chunk_size, n)
        # Try to break on a paragraph or sentence boundary near `end`.
        if end < n:
            window = text[start:end]
            for sep in ("\n\n", "\n", ". ", "。", " "):
                idx = window.rfind(sep)
                if idx > chunk_size * 0.5:  # don't break too early
                    end = start + idx + len(sep)
                    break
        chunks.append(text[start:end].strip())
        if end >= n:
            break
        start = max(end - overlap, start + 1)
    return [c for c in chunks if c]


def chunk_pages(
    pages: List[Tuple[int, str]],
    source: str,
    chunk_size: int = 800,
    overlap: int = 100,
) -> List[CitedChunk]:
    """Chunk per-page text into CitedChunks that remember their page.

    A chunk never spans two pages (so the page number is always exact).
    `chunk_index` is global across the document.
    """
    out: List[CitedChunk] = []
    idx = 0
    for page_no, page_text in pages:
        for piece in chunk_text(page_text, chunk_size, overlap):
            out.append(CitedChunk(
                text=piece, source=source, page=page_no, chunk_index=idx,
            ))
            idx += 1
    return out


def ingest_pdf(path: str, chunk_size: int = 800, overlap: int = 100) -> List[CitedChunk]:
    """Convenience: extract + chunk a PDF in one call. Source = basename."""
    pages = extract_pdf_pages(path)
    return chunk_pages(pages, os.path.basename(path), chunk_size, overlap)


# ── Obsidian / markdown vault import (Q7) ──────────────────────────────────

import re as _re

_WIKILINK_RE = _re.compile(r"\[\[([^\]|]+)(?:\|([^\]]+))?\]\]")


def _strip_wikilinks(text: str) -> Tuple[str, List[str]]:
    """Replace Obsidian `[[Note]]` / `[[Note|alias]]` with readable text and
    return the referenced note targets (for building a link graph later).

    The body keeps the human-facing display text — the alias when present,
    else the target — so the chunk reads naturally. The link graph keeps the
    real target (`extra['links']`)."""
    links: List[str] = []

    def _repl(m):
        target = m.group(1).strip()
        alias = (m.group(2) or "").strip()
        links.append(target)
        return alias or target

    return _WIKILINK_RE.sub(_repl, text), links


def import_obsidian_vault(
    vault_dir: str,
    chunk_size: int = 800,
    overlap: int = 100,
    include_subdirs: bool = True,
) -> List[CitedChunk]:
    """Walk an Obsidian vault (a folder of .md files) and return CitedChunks.

    Q7 decision: Obsidian is a first-class import because a vault is just a
    markdown folder — zero integration cost. `[[wikilinks]]` are flattened to
    their display text in the chunk body, and the referenced note names are
    kept in `extra['links']` so a future pass can build the link graph.

    `source` is the path relative to the vault root (e.g. "notes/qgt.md").
    Markdown has no pages, so `page` is None.
    """
    if not os.path.isdir(vault_dir):
        raise NotADirectoryError(f"Obsidian vault not found: {vault_dir}")

    vault_abs = os.path.abspath(vault_dir)
    out: List[CitedChunk] = []

    for root, dirs, files in os.walk(vault_abs):
        # Skip Obsidian's own config + hidden dirs.
        dirs[:] = [d for d in dirs if not d.startswith(".")]
        if not include_subdirs and root != vault_abs:
            continue
        for fname in sorted(files):
            if not fname.lower().endswith((".md", ".markdown")):
                continue
            fpath = os.path.join(root, fname)
            rel = os.path.relpath(fpath, vault_abs)
            try:
                with open(fpath, encoding="utf-8") as f:
                    raw = f.read()
            except (OSError, UnicodeDecodeError):
                continue
            clean, links = _strip_wikilinks(raw)
            idx = 0
            for piece in chunk_text(clean, chunk_size, overlap):
                out.append(CitedChunk(
                    text=piece, source=rel, page=None, chunk_index=idx,
                    extra={"links": links} if links else {},
                ))
                idx += 1
    return out
