modules/chunker.py

"""
🔪 تقسیم متن به Chunks
Smart text chunking with overlap for large documents
"""

import logging
from typing import List, Dict
from .utils import count_tokens, split_sentences, get_last_n_tokens

logger = logging.getLogger(__name__)


class TextChunker:
    """
    کلاس برای تقسیم هوشمند متن به chunks
    
    ویژگی‌ها:
    - تقسیم بر اساس تعداد tokens (نه کاراکتر)
    - همپوشانی (overlap) بین chunks برای جلوگیری از split شدن entities
    - تقسیم بر اساس جملات (نه کاراکتر) برای حفظ معنا
    - metadata کامل برای هر chunk
    """
    
    def __init__(
        self, 
        chunk_size: int = 1000,   # tokens
        overlap: int = 150         # tokens
    ):
        """
        مقداردهی اولیه chunker
        
        Args:
            chunk_size: حداکثر سایز هر chunk به tokens
            overlap: تعداد tokens همپوشانی بین chunks
            
        Examples:
            >>> chunker = TextChunker(chunk_size=1000, overlap=150)
        """
        if chunk_size <= 0:
            raise ValueError("chunk_size باید بزرگتر از 0 باشد")
        
        if overlap < 0:
            raise ValueError("overlap نمی‌تواند منفی باشد")
        
        if overlap >= chunk_size:
            raise ValueError("overlap باید کوچکتر از chunk_size باشد")
        
        self.chunk_size = chunk_size
        self.overlap = overlap
        
        logger.info(
            f"✅ TextChunker initialized: "
            f"chunk_size={chunk_size} tokens, "
            f"overlap={overlap} tokens"
        )
    
    def create_chunks(self, text: str) -> List[Dict]:
        """
        تقسیم متن به chunks با همپوشانی
        
        Args:
            text: متن ورودی
            
        Returns:
            لیستی از chunks با metadata:
            [
                {
                    "chunk_id": "chunk_01",
                    "text": "...",
                    "start_char": 0,
                    "end_char": 5000,
                    "tokens": 1000
                },
                ...
            ]
            
        Examples:
            >>> chunker = TextChunker(chunk_size=500, overlap=100)
            >>> chunks = chunker.create_chunks("متن بلند...")
            >>> len(chunks)
            3
        """
        if not text or not text.strip():
            logger.warning("⚠️ متن خالی - بازگشت لیست خالی")
            return []
        
        logger.info("🔪 شروع chunking...")
        logger.info(f"📏 طول متن: {len(text)} کاراکتر")
        
        # تقسیم به جملات
        sentences = split_sentences(text)
        
        if not sentences:
            # اگر جمله‌ای شناسایی نشد، کل متن را یک chunk در نظر بگیر
            logger.warning("⚠️ جمله‌ای شناسایی نشد - کل متن یک chunk")
            return [self._create_chunk_metadata(
                chunk_id=1,
                text=text,
                start_char=0,
                end_char=len(text),
                tokens=count_tokens(text)
            )]
        
        logger.info(f"📝 تعداد جملات: {len(sentences)}")
        
        chunks = []
        current_chunk = ""
        current_tokens = 0
        current_start_char = 0
        
        for sentence in sentences:
            sentence_tokens = count_tokens(sentence)
            
            # آیا اضافه کردن این جمله chunk را بیش از حد بزرگ می‌کند؟
            if current_tokens + sentence_tokens > self.chunk_size and current_chunk:
                # ذخیره chunk فعلی
                chunk_text = current_chunk.strip()
                chunks.append(self._create_chunk_metadata(
                    chunk_id=len(chunks) + 1,
                    text=chunk_text,
                    start_char=current_start_char,
                    end_char=current_start_char + len(chunk_text),
                    tokens=current_tokens
                ))
                
                # شروع chunk جدید با overlap
                overlap_text = get_last_n_tokens(current_chunk, self.overlap)
                
                # محاسبه موقعیت شروع جدید
                current_start_char = current_start_char + len(current_chunk) - len(overlap_text)
                
                # شروع chunk جدید
                current_chunk = overlap_text + " " + sentence
                current_tokens = count_tokens(current_chunk)
            else:
                # اضافه کردن جمله به chunk فعلی
                if current_chunk:
                    current_chunk += " "
                current_chunk += sentence
                current_tokens += sentence_tokens
        
        # ذخیره chunk آخر
        if current_chunk:
            chunk_text = current_chunk.strip()
            chunks.append(self._create_chunk_metadata(
                chunk_id=len(chunks) + 1,
                text=chunk_text,
                start_char=current_start_char,
                end_char=current_start_char + len(chunk_text),
                tokens=current_tokens
            ))
        
        # لاگ نتیجه
        logger.info(f"✅ تقسیم به {len(chunks)} chunk انجام شد")
        
        for chunk in chunks:
            logger.info(
                f"  {chunk['chunk_id']}: "
                f"{chunk['tokens']} tokens, "
                f"{len(chunk['text'])} chars"
            )
        
        return chunks
    
    def _create_chunk_metadata(
        self, 
        chunk_id: int, 
        text: str,
        start_char: int,
        end_char: int,
        tokens: int
    ) -> Dict:
        """
        ساخت metadata برای یک chunk
        
        Args:
            chunk_id: شماره chunk
            text: متن chunk
            start_char: موقعیت شروع در متن اصلی (کاراکتر)
            end_char: موقعیت پایان در متن اصلی (کاراکتر)
            tokens: تعداد tokens
            
        Returns:
            دیکشنری حاوی metadata
        """
        return {
            "chunk_id": f"chunk_{chunk_id:02d}",
            "text": text,
            "start_char": start_char,
            "end_char": end_char,
            "tokens": tokens,
            "length": len(text)
        }
    
    def validate_chunks(self, chunks: List[Dict]) -> bool:
        """
        اعتبارسنجی chunks
        
        بررسی می‌کند که:
        - تمام chunks معتبر هستند
        - هیچ chunk خالی نیست
        - overlaps درست هستند
        
        Args:
            chunks: لیست chunks
            
        Returns:
            True اگر همه چیز درست باشد
        """
        if not chunks:
            return True
        
        for chunk in chunks:
            # بررسی فیلدهای ضروری
            required_fields = ["chunk_id", "text", "start_char", "end_char", "tokens"]
            for field in required_fields:
                if field not in chunk:
                    logger.error(f"❌ Chunk {chunk.get('chunk_id', '?')} فیلد {field} ندارد")
                    return False
            
            # بررسی خالی نبودن
            if not chunk["text"].strip():
                logger.error(f"❌ Chunk {chunk['chunk_id']} خالی است")
                return False
            
            # بررسی tokens
            if chunk["tokens"] <= 0:
                logger.error(f"❌ Chunk {chunk['chunk_id']} تعداد tokens نامعتبر دارد")
                return False
        
        logger.info("✅ تمام chunks معتبر هستند")
        return True


# ✅ تست‌های سریع
if __name__ == "__main__":
    print("=" * 60)
    print("🧪 Testing Chunker Module")
    print("=" * 60)
    
    # تست 1: Chunking ساده
    print("\n📊 Test 1: Simple Chunking")
    
    test_text = """
    شرکت پارس در سال گذشته فروش 50 میلیارد ریال داشت. این رقم نسبت به سال قبل رشد 15 درصدی دارد.
    علی احمدی مدیرعامل شرکت است. شرکت صبا نیز در همین بازار فعالیت می‌کند.
    شرکت صبا فروش 30 میلیارد ریال داشت. مریم کریمی مدیر مالی این شرکت است.
    همکاری بین دو شرکت در دستور کار است. قرارداد به ارزش 20 میلیارد ریال است.
    سرمایه‌گذاری 10 میلیارد ریال انجام خواهد شد. بازده پیش‌بینی شده 25 درصد است.
    """
    
    chunker = TextChunker(chunk_size=100, overlap=20)
    chunks = chunker.create_chunks(test_text)
    
    print(f"  متن اصلی: {len(test_text)} کاراکتر")
    print(f"  تعداد chunks: {len(chunks)}")
    
    for chunk in chunks:
        print(f"\n  {chunk['chunk_id']}:")
        print(f"    Tokens: {chunk['tokens']}")
        print(f"    Length: {chunk['length']} chars")
        print(f"    Position: [{chunk['start_char']}, {chunk['end_char']}]")
        print(f"    Preview: {chunk['text'][:80]}...")
    
    # تست 2: Validation
    print("\n📊 Test 2: Chunk Validation")
    is_valid = chunker.validate_chunks(chunks)
    print(f"  Valid: {is_valid}")
    
    # تست 3: متن خیلی کوتاه
    print("\n📊 Test 3: Very Short Text")
    short_text = "این یک متن کوتاه است."
    short_chunks = chunker.create_chunks(short_text)
    print(f"  تعداد chunks: {len(short_chunks)}")
    
    # تست 4: متن خالی
    print("\n📊 Test 4: Empty Text")
    empty_chunks = chunker.create_chunks("")
    print(f"  تعداد chunks: {len(empty_chunks)}")
    
    print("\n" + "=" * 60)
    print("✅ All tests completed!")
    print("=" * 60)