import hashlib
import re
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path

from techdocs_graphrag.models import Chunk

SUPPORTED_SUFFIXES = frozenset({".md", ".markdown", ".txt"})
PARAGRAPH_SEPARATOR = "\n\n"


@dataclass(frozen=True, slots=True)
class SourceDocument:
    name: str
    text: str


def load_documents(root: Path) -> list[SourceDocument]:
    return [
        SourceDocument(name=path.relative_to(root).as_posix(), text=path.read_text(encoding="utf-8"))
        for path in sorted(root.rglob("*"))
        if path.is_file() and path.suffix.lower() in SUPPORTED_SUFFIXES
    ]


class ParagraphSplitter:
    def __init__(self, chunk_size: int, chunk_overlap: int) -> None:
        if not 0 <= chunk_overlap < chunk_size:
            raise ValueError("chunk_overlap must be in [0, chunk_size)")
        self._chunk_size = chunk_size
        self._chunk_overlap = chunk_overlap

    def split(self, text: str) -> list[str]:
        chunks: list[str] = []
        buffer = ""
        for paragraph in self._paragraphs(text):
            candidate = self._join(buffer, paragraph)
            if len(candidate) <= self._chunk_size:
                buffer = candidate
                continue
            if buffer:
                chunks.append(buffer)
                candidate = self._join(self._overlap_of(buffer), paragraph)
            while len(candidate) > self._chunk_size:
                chunks.append(candidate[: self._chunk_size])
                candidate = candidate[self._chunk_size - self._chunk_overlap :]
            buffer = candidate
        if buffer:
            chunks.append(buffer)
        return chunks

    @staticmethod
    def _paragraphs(text: str) -> Iterable[str]:
        return (block.strip() for block in re.split(r"\n\s*\n", text) if block.strip())

    @staticmethod
    def _join(head: str, tail: str) -> str:
        return f"{head}{PARAGRAPH_SEPARATOR}{tail}" if head else tail

    def _overlap_of(self, text: str) -> str:
        if not self._chunk_overlap or len(text) <= self._chunk_overlap:
            return ""
        tail = text[-self._chunk_overlap :]
        return tail.partition(" ")[2] or tail


def chunk_id(document: str, position: int, text: str) -> str:
    return hashlib.sha256(f"{document}\x00{position}\x00{text}".encode()).hexdigest()[:16]


def chunk_documents(documents: Iterable[SourceDocument], splitter: ParagraphSplitter) -> list[Chunk]:
    return [
        Chunk(id=chunk_id(document.name, position, text), document=document.name, position=position, text=text)
        for document in documents
        for position, text in enumerate(splitter.split(document.text))
    ]