from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from langchain_text_splitters import RecursiveCharacterTextSplitter

from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger

logger = get_logger(__name__)


@dataclass
class TextChunk:
    content: str
    chunk_index: int
    start_char: int
    end_char: int
    metadata: dict[str, Any]


class SemanticChunker:
    def __init__(self):
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=settings.chunking.chunk_size,
            chunk_overlap=settings.chunking.chunk_overlap,
            separators=settings.chunking.separators,
            keep_separator=settings.chunking.keep_separator,
            length_function=len,
        )

    def chunk(self, text: str, base_metadata: dict[str, Any] | None = None) -> list[TextChunk]:
        if not text or not text.strip():
            return []

        base_metadata = base_metadata or {}
        docs = self.splitter.create_documents([text], metadatas=[base_metadata])

        chunks = []
        for i, doc in enumerate(docs):
            start = text.find(doc.page_content)
            if start == -1:
                start = 0
            chunks.append(
                TextChunk(
                    content=doc.page_content,
                    chunk_index=i,
                    start_char=start,
                    end_char=start + len(doc.page_content),
                    metadata={**doc.metadata, "chunk_index": i, "total_chunks": len(docs)},
                )
            )

        logger.debug("chunked_text", chunks=len(chunks), original_length=len(text))
        return chunks


def chunk_text(text: str, metadata: dict[str, Any] | None = None) -> list[TextChunk]:
    chunker = SemanticChunker()
    return chunker.chunk(text, metadata)