"""Lightweight workbook decomposition and retrieval for Excel Logic Wizard."""
from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Dict, Iterable, List, Tuple


@dataclass
class WorkbookChunk:
    chunk_id: str
    chunk_type: str
    text: str
    metadata: Dict[str, str]


class ExcelWorkbookRAG:
    """Simple lexical retriever over workbook decomposition chunks."""

    def __init__(self):
        self._chunks: List[WorkbookChunk] = []

    @staticmethod
    def _tokenize(text: str) -> List[str]:
        return re.findall(r"[a-z0-9_]+", (text or "").lower())

    def add_chunks(self, chunks: Iterable[WorkbookChunk]) -> None:
        self._chunks.extend(list(chunks))

    def chunk_count(self) -> int:
        return len(self._chunks)

    def retrieve(self, query: str, top_k: int = 8) -> List[Tuple[WorkbookChunk, float]]:
        q_tokens = set(self._tokenize(query))
        if not q_tokens:
            return []

        scored: List[Tuple[WorkbookChunk, float]] = []
        for c in self._chunks:
            c_tokens = set(self._tokenize(c.text))
            if not c_tokens:
                continue
            inter = len(q_tokens & c_tokens)
            union = len(q_tokens | c_tokens)
            if union == 0:
                continue
            score = inter / union
            if score > 0:
                scored.append((c, score))

        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

