"""
TF-IDF based context extractor — Python port of src/context-extractor.js

Given a large document + a query, returns only the relevant chunks within a
token budget. No external dependencies. Runs in milliseconds.

Usage::

    from mdmin import extract

    result = extract(markdown_text, "how does auth work", max_tokens=2000)
    print(result.text)          # relevant chunks only
    print(result.stats.reduction)  # % reduction
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass, field
from typing import Optional

# ─── Stop words ──────────────────────────────────────────────────────────────

_STOP_WORDS: frozenset[str] = frozenset({
    'a', 'an', 'the', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for',
    'of', 'with', 'by', 'from', 'is', 'it', 'its', 'be', 'as', 'was',
    'are', 'has', 'had', 'have', 'not', 'this', 'that', 'can', 'all',
    'if', 'do', 'use', 'also', 'new', 'get', 'set', 'you', 'your', 'our',
    'we', 'they', 'he', 'she', 'his', 'her', 'us', 'any', 'each', 'into',
    'see', 'way', 'may', 'too', 'so', 'no', 'up', 'out', 'now', 'then',
    'than', 'when', 'how', 'what', 'which', 'who', 'only', 'just', 'about',
    'other', 'after', 'before', 'over', 'such', 'own', 'same', 'per',
    'both', 'must', 'some', 'one', 'two', 'three', 'first', 'second', 'add',
    'note', 'make', 'like', 'need', 'take', 'more', 'very', 'been', 'being',
    'through', 'time', 'turn', 'well', 'work', 'will', 'would', 'there',
    'their', 'should', 'does', 'were', 'much', 'between', 'these', 'those',
})

# ─── Token estimation ─────────────────────────────────────────────────────────


def _estimate_tokens(text: str) -> int:
    if not text:
        return 0
    tokens = 0
    for chunk in re.findall(r'\S+|\s+', text):
        if re.match(r'^\s+$', chunk):
            tokens += chunk.count('\n')
        else:
            length = len(chunk)
            if length <= 6:
                tokens += 1
            elif length <= 10:
                tokens += 2
            elif length <= 15:
                tokens += 3
            else:
                tokens += math.ceil(length / 4)
    return round(tokens)


# ─── TF-IDF ───────────────────────────────────────────────────────────────────


class _TFIDF:
    def __init__(self) -> None:
        self._idf: dict[str, float] = {}
        self._term_freqs: list[dict[str, int]] = []

    def _tokenize(self, text: str) -> list[str]:
        return [
            w for w in re.sub(r'[^a-z0-9\s]', ' ', text.lower()).split()
            if len(w) > 2 and w not in _STOP_WORDS
        ]

    def fit(self, documents: list[str]) -> None:
        doc_count = len(documents)
        df: dict[str, int] = {}

        self._term_freqs = []
        for doc in documents:
            terms = self._tokenize(doc)
            tf: dict[str, int] = {}
            seen: set[str] = set()
            for term in terms:
                tf[term] = tf.get(term, 0) + 1
                if term not in seen:
                    df[term] = df.get(term, 0) + 1
                    seen.add(term)
            self._term_freqs.append(tf)

        self._idf = {
            term: math.log((doc_count + 1) / (count + 1)) + 1
            for term, count in df.items()
        }

    def vectorize(self, text: str) -> dict[str, float]:
        terms = self._tokenize(text)
        if not terms:
            return {}
        tf: dict[str, int] = {}
        for term in terms:
            tf[term] = tf.get(term, 0) + 1
        max_tf = max(tf.values()) if tf else 1
        return {
            term: (count / max_tf) * self._idf.get(term, 1.0)
            for term, count in tf.items()
        }

    def similarity(self, vec_a: dict[str, float], vec_b: dict[str, float]) -> float:
        if not vec_a or not vec_b:
            return 0.0
        dot = sum(vec_a[t] * vec_b[t] for t in vec_a if t in vec_b)
        mag_a = math.sqrt(sum(v * v for v in vec_a.values()))
        mag_b = math.sqrt(sum(v * v for v in vec_b.values()))
        if mag_a == 0 or mag_b == 0:
            return 0.0
        return dot / (mag_a * mag_b)


# ─── Data types ───────────────────────────────────────────────────────────────


@dataclass
class TopScore:
    preview: str
    score: float


@dataclass
class ExtractStats:
    total_doc_tokens: int
    extracted_tokens: int
    chunks_total: int
    chunks_extracted: int
    reduction: float
    top_scores: list[TopScore] = field(default_factory=list)


@dataclass
class ExtractResult:
    text: str
    stats: ExtractStats


# ─── ContextExtractor ─────────────────────────────────────────────────────────


class ContextExtractor:
    """
    Index a markdown document then extract query-relevant chunks within a
    token budget using TF-IDF cosine similarity.

    Example::

        extractor = ContextExtractor()
        extractor.index(markdown_text)
        result = extractor.extract("how does auth work", max_tokens=2000)
        print(result.text)
    """

    def __init__(self) -> None:
        self._chunks: list[dict] = []
        self._tfidf = _TFIDF()
        self._chunk_vectors: list[dict[str, float]] = []
        self._indexed = False

    def index(self, markdown: str, max_chunk_tokens: int = 300) -> dict:
        """Index a document. Must be called before extract()."""
        self._chunks = self._chunk_document(markdown, max_chunk_tokens)
        self._tfidf.fit([c['text'] for c in self._chunks])
        self._chunk_vectors = [self._tfidf.vectorize(c['text']) for c in self._chunks]
        self._indexed = True
        return {
            'chunksCreated': len(self._chunks),
            'totalTokens': sum(c['tokens'] for c in self._chunks),
        }

    def extract(self, query: str, max_tokens: int = 2000,
                min_relevance: float = 0.05, include_context: bool = True) -> ExtractResult:
        """Extract chunks relevant to *query* within *max_tokens* budget."""
        if not self._indexed:
            raise RuntimeError('Call index() before extract()')

        qv = self._tfidf.vectorize(query)

        scored = sorted(
            [
                {'chunk': c, 'index': i,
                 'score': self._tfidf.similarity(qv, self._chunk_vectors[i])}
                for i, c in enumerate(self._chunks)
            ],
            key=lambda x: x['score'],
            reverse=True,
        )

        selected: set[int] = set()
        token_count = 0

        for item in scored:
            if item['score'] < min_relevance:
                break
            chunk = item['chunk']
            idx = item['index']
            if token_count + chunk['tokens'] > max_tokens:
                continue
            selected.add(idx)
            token_count += chunk['tokens']

            if include_context and idx > 0 and (idx - 1) not in selected:
                prev = self._chunks[idx - 1]
                if token_count + prev['tokens'] <= max_tokens:
                    selected.add(idx - 1)
                    token_count += prev['tokens']

        indices = sorted(selected)
        text = '\n\n'.join(self._chunks[i]['text'] for i in indices)

        total_tokens = sum(c['tokens'] for c in self._chunks)
        extracted_tokens = _estimate_tokens(text)
        reduction = round(((total_tokens - extracted_tokens) / total_tokens) * 1000) / 10 if total_tokens else 0.0

        top_scores = [
            TopScore(
                preview=s['chunk']['text'][:80] + '...',
                score=round(s['score'] * 1000) / 1000,
            )
            for s in scored[:5]
        ]

        return ExtractResult(
            text=text,
            stats=ExtractStats(
                total_doc_tokens=total_tokens,
                extracted_tokens=extracted_tokens,
                chunks_total=len(self._chunks),
                chunks_extracted=len(selected),
                reduction=reduction,
                top_scores=top_scores,
            ),
        )

    def score_chunks(self, query: str) -> list[dict]:
        """Return all chunks scored against *query*, sorted descending by score.

        Useful for multi-doc extraction — score chunks from multiple documents
        then select globally within a combined token budget.

        Returns list of dicts: ``{'chunk': dict, 'index': int, 'score': float}``
        """
        if not self._indexed:
            raise RuntimeError('Call index() before score_chunks()')
        qv = self._tfidf.vectorize(query)
        return sorted(
            [
                {'chunk': c, 'index': i,
                 'score': self._tfidf.similarity(qv, self._chunk_vectors[i])}
                for i, c in enumerate(self._chunks)
            ],
            key=lambda x: x['score'],
            reverse=True,
        )

    # ── Private helpers ───────────────────────────────────────────────────────

    def _chunk_document(self, markdown: str, max_tokens: int) -> list[dict]:
        chunks: list[dict] = []
        position = 0

        sections = re.split(r'(?=^#{1,3}\s)', markdown, flags=re.MULTILINE)

        for section in sections:
            section_tokens = _estimate_tokens(section)
            if section_tokens <= max_tokens:
                if section.strip():
                    chunks.append({
                        'text': section.strip(),
                        'tokens': section_tokens,
                        'position': position,
                        'header': self._extract_header(section),
                    })
                    position += 1
            else:
                paragraphs = re.split(r'\n\n+', section)
                buffer = ''
                header_seen = False

                for para in paragraphs:
                    header = self._extract_header(section) if not header_seen else None
                    if header:
                        header_seen = True

                    candidate = buffer + '\n\n' + para if buffer else para
                    if _estimate_tokens(candidate) > max_tokens and buffer:
                        chunks.append({
                            'text': buffer.strip(),
                            'tokens': _estimate_tokens(buffer),
                            'position': position,
                            'header': header,
                        })
                        position += 1
                        buffer = ''
                    buffer = candidate

                if buffer.strip():
                    chunks.append({
                        'text': buffer.strip(),
                        'tokens': _estimate_tokens(buffer),
                        'position': position,
                        'header': self._extract_header(section),
                    })
                    position += 1

        return chunks

    def _extract_header(self, text: str) -> Optional[dict]:
        m = re.search(r'^(#{1,6})\s+(.+)$', text, re.MULTILINE)
        return {'level': len(m.group(1)), 'text': m.group(2)} if m else None


# ─── Convenience function ─────────────────────────────────────────────────────


def quick_extract(markdown: str, query: str, options: Optional[dict] = None) -> ExtractResult:
    """One-call extraction helper."""
    opts = options or {}
    extractor = ContextExtractor()
    extractor.index(markdown)
    return extractor.extract(
        query,
        max_tokens=opts.get('maxTokens', 2000),
        min_relevance=opts.get('minRelevance', 0.05),
    )
