from dataclasses import dataclass
from typing import Optional
from config import settings
from logger import get_logger

log = get_logger("chunker")


@dataclass
class Chunk:
    text: str
    index: int
    source: str
    page: Optional[int] = None
    sheet: Optional[str] = None
    metadata: Optional[dict] = None


class TextChunker:
    def __init__(
        self,
        chunk_size: int = None,
        overlap: int = None
    ):
        self.chunk_size = chunk_size or settings.CHUNK_SIZE
        self.overlap = overlap or settings.CHUNK_OVERLAP
        log.debug(f"TextChunker: size={self.chunk_size}, overlap={self.overlap}")
    
    def chunk_text(
        self,
        text: str,
        source: str,
        page: Optional[int] = None,
        sheet: Optional[str] = None,
        metadata: Optional[dict] = None
    ) -> list[Chunk]:
        words = text.split()
        log.debug(f"Chunking text: {len(words)} words, source={source}, page={page}, sheet={sheet}")
        
        if len(words) <= self.chunk_size:
            log.debug(f"Text fits in single chunk")
            return [Chunk(
                text=text,
                index=0,
                source=source,
                page=page,
                sheet=sheet,
                metadata=metadata
            )]
        
        chunks = []
        start = 0
        idx = 0
        
        while start < len(words):
            end = start + self.chunk_size
            chunk_words = words[start:end]
            chunk_text = " ".join(chunk_words)
            
            chunks.append(Chunk(
                text=chunk_text,
                index=idx,
                source=source,
                page=page,
                sheet=sheet,
                metadata=metadata
            ))
            
            log.debug(f"Chunk {idx}: {len(chunk_words)} words, pos {start}-{end}")
            start = end - self.overlap
            idx += 1
        
        log.debug(f"Created {len(chunks)} chunks")
        return chunks
    
    def chunk_document(self, doc: dict) -> list[Chunk]:
        log.info(f"Chunking document: {doc.get('source')}")
        all_chunks = []
        
        if "pages" in doc:
            log.debug(f"Document has {len(doc['pages'])} pages/sheets")
            for page_data in doc["pages"]:
                page_chunks = self.chunk_text(
                    text=page_data["text"],
                    source=doc["source"],
                    page=page_data.get("page"),
                    sheet=page_data.get("sheet"),
                    metadata={"type": doc["type"]}
                )
                all_chunks.extend(page_chunks)
                if page_data.get("page"):
                    log.debug(f"Page {page_data['page']}: {len(page_chunks)} chunks")
                elif page_data.get("sheet"):
                    log.debug(f"Sheet '{page_data['sheet']}': {len(page_chunks)} chunks")
        else:
            log.debug(f"Single page document")
            all_chunks = self.chunk_text(
                text=doc["full_text"],
                source=doc["source"],
                metadata={"type": doc["type"]}
            )
        
        log.info(f"Document chunked: {len(all_chunks)} total chunks")
        return all_chunks


def chunk_document(doc: dict) -> list[Chunk]:
    chunker = TextChunker()
    return chunker.chunk_document(doc)