import tiktoken
from .config import CONFIG

_enc = tiktoken.get_encoding("cl100k_base")


def count_tokens(text):
    return len(_enc.encode(text))


def chunk_text(text, chunk_size=None, overlap=None):
    chunk_size = chunk_size or CONFIG.chunk_size_tokens
    overlap = overlap or CONFIG.chunk_overlap_tokens
    
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    
    chunks = []
    current_chunk = []
    current_tokens = 0
    
    for para in paragraphs:
        para_tokens = count_tokens(para)
        
        if para_tokens > chunk_size:
            sentences = para.replace("\n", " ").split(". ")
            for sent in sentences:
                if not sent.strip():
                    continue
                sent_tokens = count_tokens(sent)
                if current_tokens + sent_tokens > chunk_size and current_chunk:
                    chunks.append("\n\n".join(current_chunk))
                    current_chunk = current_chunk[-1:] if current_chunk else []
                    current_tokens = count_tokens("\n\n".join(current_chunk)) if current_chunk else 0
                current_chunk.append(sent)
                current_tokens += sent_tokens
        else:
            if current_tokens + para_tokens > chunk_size and current_chunk:
                chunks.append("\n\n".join(current_chunk))
                overlap_tokens = 0
                overlap_paras = []
                for p in reversed(current_chunk):
                    pt = count_tokens(p)
                    if overlap_tokens + pt <= overlap:
                        overlap_paras.insert(0, p)
                        overlap_tokens += pt
                    else:
                        break
                current_chunk = overlap_paras
                current_tokens = overlap_tokens
            current_chunk.append(para)
            current_tokens += para_tokens
    
    if current_chunk:
        chunks.append("\n\n".join(current_chunk))
    
    return chunks