"""Chunked multi-pass extraction for large documents.

When a source document exceeds the single-pass context limit (~10K chars),
this module splits it into chunks, runs parallel LLM extraction on each,
then merges the results into a single consolidated summary that captures
the ENTIRE document — no content is silently dropped.

Flow:
  full_text  →  chunk (10K segments)
             →  parallel LLM extract per chunk
             →  LLM merge all chunk extractions
             →  consolidated text (fed to normal classify pipeline)

Design:
- Chunks split on paragraph boundaries (never mid-sentence)
- Each chunk gets positional context ("Section 3 of 12")
- Merge pass deduplicates concepts, resolves cross-references
- The caller (KBService.ingest) gets back a single string that
  replaces the old truncated source_text
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from agntspace.llm.base import LLMProvider

log = logging.getLogger(__name__)


# ── Phase 2 (Rule 12) chunking config ──────────────────────────────────────
#
# Strategy lives in ``defaults/skills/core/kb-ingest/config/chunking.yaml``.
# main.py loads it at startup via ``set_chunking_config(...)``. The Python
# code here is a thin engine that uses whatever the config says.
#
# The pattern mirrors ``integrations/browser.py``'s ``ScrapingConfig`` and
# ``set_scraping_config()`` -- module-level dataclass + setter so the
# call sites stay free of skill-loader injection.

@dataclass
class ChunkingConfig:
    """Per-document multi-pass chunking strategy.

    Loaded from kb-ingest/config/chunking.yaml at startup. The fields
    here are the only knobs the chunker exposes; tuning happens by
    editing the YAML, never by editing this dataclass.
    """
    # Defaults match the pre-Phase-2 hardcoded constants exactly so
    # the migration produces zero behavior change. They are also the
    # documented fallbacks for when the YAML is missing entirely.
    chunk_threshold_chars: int = 10_000
    chunk_size_chars: int = 10_000
    chunk_overlap_chars: int = 500
    max_chunks: int = 50


_CONFIG = ChunkingConfig()  # arch:deterministic -- module-level config holder; the actual values come from set_chunking_config(skill YAML) at startup, see Phase 2 docstring


def set_chunking_config(cfg: dict | None) -> None:
    """Install chunking strategy from the kb-ingest skill config.

    Called once at server startup from ``main.py`` after the SkillLoader
    has loaded ``chunking.yaml``. Empty / missing config keeps the
    pre-Phase-2 defaults so degraded mode behaves identically to the
    pre-migration code.
    """
    global _CONFIG
    if not cfg or not isinstance(cfg, dict):
        return
    _CONFIG = ChunkingConfig(
        chunk_threshold_chars=int(cfg.get("chunk_threshold_chars", _CONFIG.chunk_threshold_chars)),
        chunk_size_chars=int(cfg.get("chunk_size_chars", _CONFIG.chunk_size_chars)),
        chunk_overlap_chars=int(cfg.get("chunk_overlap_chars", _CONFIG.chunk_overlap_chars)),
        max_chunks=int(cfg.get("max_chunks", _CONFIG.max_chunks)),
    )
    log.info(
        "Chunking config: threshold=%d size=%d overlap=%d max_chunks=%d",
        _CONFIG.chunk_threshold_chars,
        _CONFIG.chunk_size_chars,
        _CONFIG.chunk_overlap_chars,
        _CONFIG.max_chunks,
    )


def get_chunking_config() -> ChunkingConfig:
    """Return the current ChunkingConfig (used by tests + diagnostics)."""
    return _CONFIG


def needs_chunking(text: str) -> bool:
    """Check if text is large enough to require multi-pass extraction."""
    return len(text) > _CONFIG.chunk_threshold_chars


def split_into_chunks(text: str) -> list[str]:
    """Split text into overlapping chunks on paragraph boundaries.

    Returns a list of text segments, each roughly chunk_size_chars,
    split at paragraph boundaries (\n\n) to avoid cutting mid-thought.
    Adjacent chunks overlap by chunk_overlap_chars for context continuity.
    """
    chunk_size = _CONFIG.chunk_size_chars
    chunk_overlap = _CONFIG.chunk_overlap_chars
    max_chunks = _CONFIG.max_chunks

    if len(text) <= chunk_size:
        return [text]

    chunks: list[str] = []
    start = 0
    text_len = len(text)

    while start < text_len:
        end = min(start + chunk_size, text_len)

        # If not at the end, find a paragraph boundary to split on
        if end < text_len:
            # Look for \n\n near the end of the chunk (search last 20%)
            search_start = end - chunk_size // 5
            split_pos = text.rfind("\n\n", search_start, end)
            if split_pos > start:
                end = split_pos + 2  # include the \n\n

        chunks.append(text[start:end])

        # Advance with overlap
        start = end - chunk_overlap if end < text_len else text_len

        if len(chunks) >= max_chunks:
            remaining = text_len - start
            if remaining > 0:
                log.warning("Document too large: truncated at %d chunks (%d chars remaining)",
                            max_chunks, remaining)
            break

    return chunks


_CHUNK_EXTRACT_PROMPT = """You are extracting knowledge from one section of a larger document.

This is section {chunk_num} of {total_chunks}.
Source file: {filename}
Source type: {source_type}

Extract ALL important information from this section:
- Key facts, claims, and findings
- People, organizations, and entities mentioned
- Concepts, methodologies, and frameworks
- Decisions, conclusions, and recommendations
- Data points, statistics, and metrics
- Relationships between entities

Be thorough — information not captured here is LOST. Write a dense, factual summary.
Do NOT add commentary or analysis. Just extract what's in the text.

Produce your extraction as a structured summary with bullet points grouped by topic."""

_MERGE_PROMPT = """You are consolidating extracted knowledge from {total_chunks} sections of one document.

Source file: {filename}
Source type: {source_type}

Below are the section-by-section extractions. Your job:
1. Merge into ONE coherent, comprehensive summary
2. Deduplicate — same fact mentioned in multiple sections appears once
3. Resolve cross-references (e.g., "the method described above" → name it)
4. Preserve ALL unique facts — do not drop information
5. Organize by topic, not by section number
6. Keep entity names, dates, and specifics — no vague paraphrasing

The output will be used for wiki article generation, so be factual and dense.
Every claim must trace to the source document.

SECTION EXTRACTIONS:
{extractions}"""


async def chunked_extract(
    text: str,
    filename: str,
    source_type: str,
    llm: LLMProvider,
    *,
    model: str | None = None,
) -> str:
    """Extract knowledge from a large document via chunked multi-pass.

    1. Split text into chunks
    2. Run parallel LLM extraction on each chunk
    3. Merge all extractions into a single consolidated summary

    Returns the consolidated summary text (replaces the old truncated source_text).
    """
    from agntspace.llm.base import Message

    chunks = split_into_chunks(text)
    total = len(chunks)
    log.info("Chunked extraction: %s → %d chunks (%d chars total)",
             filename, total, len(text))

    if total == 1:
        # Single chunk — no merge needed, just return the text directly
        return text

    # Phase 1: Parallel extraction from each chunk
    async def extract_chunk(i: int, chunk: str) -> str:
        system = _CHUNK_EXTRACT_PROMPT.format(
            chunk_num=i + 1,
            total_chunks=total,
            filename=filename,
            source_type=source_type,
        )
        try:
            resp = await llm.complete(
                messages=[Message(role="user", content=chunk)],
                system=system,
                max_tokens=3000,
                temperature=0.2,
                model=model,
            )
            return resp.content
        except Exception:
            log.exception("Chunk %d/%d extraction failed for %s", i + 1, total, filename)
            # Fallback: return a truncated version of the chunk itself
            return chunk[:2000]

    # Run all chunk extractions concurrently (with semaphore to limit parallelism)
    semaphore = asyncio.Semaphore(5)  # max 5 concurrent LLM calls

    async def extract_with_limit(i: int, chunk: str) -> str:
        async with semaphore:
            return await extract_chunk(i, chunk)

    tasks = [extract_with_limit(i, chunk) for i, chunk in enumerate(chunks)]
    extractions = await asyncio.gather(*tasks)

    log.info("Chunked extraction: %d/%d chunks extracted for %s",
             sum(1 for e in extractions if e), total, filename)

    # Phase 2: Merge all extractions into one consolidated summary
    combined = "\n\n".join(
        f"--- Section {i + 1}/{total} ---\n{ext}"
        for i, ext in enumerate(extractions)
        if ext.strip()
    )

    # If merged extractions are small enough, return directly
    if len(combined) <= _CONFIG.chunk_size_chars:
        return combined

    merge_system = _MERGE_PROMPT.format(
        total_chunks=total,
        filename=filename,
        source_type=source_type,
        extractions="",  # goes as user message for length
    )

    try:
        resp = await llm.complete(
            messages=[Message(role="user", content=combined)],
            system=merge_system,
            max_tokens=4000,
            temperature=0.2,
            model=model,
        )
        merged = resp.content
        log.info("Chunked extraction: merge complete for %s (%d chars → %d chars)",
                 filename, len(text), len(merged))
        return merged
    except Exception:
        log.exception("Merge pass failed for %s — returning concatenated extractions", filename)
        # Fallback: return concatenated extractions (may be long, but better than losing data)
        return combined
