#!/usr/bin/env python3
"""
Chunking Optimizer - Analyzes document corpus and recommends optimal chunking strategy.

This script analyzes a collection of text/markdown documents and evaluates different
chunking strategies to recommend the optimal approach for the given corpus.

Strategies tested:
- Fixed-size chunking (character and token-based) with overlap
- Sentence-based chunking
- Paragraph-based chunking  
- Semantic chunking (heading-aware)

Metrics measured:
- Chunk size distribution (mean, std, min, max)
- Semantic coherence (topic continuity heuristic)
- Boundary quality (sentence break analysis)

No external dependencies - uses only Python standard library.
"""

import argparse
import json
import os
import re
import statistics
from collections import Counter, defaultdict
from math import log, sqrt
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any


class DocumentCorpus:
    """Handles loading and preprocessing of document corpus."""
    
    def __init__(self, directory: str, extensions: List[str] = None):
        self.directory = Path(directory)
        self.extensions = extensions or ['.txt', '.md', '.markdown']
        self.documents = []
        self._load_documents()
    
    def _load_documents(self):
        """Load all text documents from directory."""
        if not self.directory.exists():
            raise FileNotFoundError(f"Directory not found: {self.directory}")
        
        for file_path in self.directory.rglob('*'):
            if file_path.is_file() and file_path.suffix.lower() in self.extensions:
                try:
                    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                        content = f.read()
                    if content.strip():  # Only include non-empty files
                        self.documents.append({
                            'path': str(file_path),
                            'content': content,
                            'size': len(content)
                        })
                except Exception as e:
                    print(f"Warning: Could not read {file_path}: {e}")
        
        if not self.documents:
            raise ValueError(f"No valid documents found in {self.directory}")
        
        print(f"Loaded {len(self.documents)} documents totaling {sum(d['size'] for d in self.documents):,} characters")


class ChunkingStrategy:
    """Base class for chunking strategies."""
    
    def __init__(self, name: str, config: Dict[str, Any]):
        self.name = name
        self.config = config
    
    def chunk(self, text: str) -> List[Dict[str, Any]]:
        """Split text into chunks. Returns list of chunk dictionaries."""
        raise NotImplementedError


class FixedSizeChunker(ChunkingStrategy):
    """Fixed-size chunking with optional overlap."""
    
    def __init__(self, chunk_size: int = 1000, overlap: int = 100, unit: str = 'char'):
        config = {'chunk_size': chunk_size, 'overlap': overlap, 'unit': unit}
        super().__init__(f'fixed_size_{unit}', config)
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.unit = unit
    
    def chunk(self, text: str) -> List[Dict[str, Any]]:
        chunks = []
        if self.unit == 'char':
            return self._chunk_by_chars(text)
        else:  # word-based approximation
            words = text.split()
            return self._chunk_by_words(words)
    
    def _chunk_by_chars(self, text: str) -> List[Dict[str, Any]]:
        chunks = []
        start = 0
        chunk_id = 0
        
        while start < len(text):
            end = min(start + self.chunk_size, len(text))
            chunk_text = text[start:end]
            
            chunks.append({
                'id': chunk_id,
                'text': chunk_text,
                'start': start,
                'end': end,
                'size': len(chunk_text)
            })
            
            start = max(start + self.chunk_size - self.overlap, start + 1)
            chunk_id += 1
            
            if start >= len(text):
                break
        
        return chunks
    
    def _chunk_by_words(self, words: List[str]) -> List[Dict[str, Any]]:
        chunks = []
        start = 0
        chunk_id = 0
        
        while start < len(words):
            end = min(start + self.chunk_size, len(words))
            chunk_words = words[start:end]
            chunk_text = ' '.join(chunk_words)
            
            chunks.append({
                'id': chunk_id,
                'text': chunk_text,
                'start': start,
                'end': end,
                'size': len(chunk_text)
            })
            
            start = max(start + self.chunk_size - self.overlap, start + 1)
            chunk_id += 1
            
            if start >= len(words):
                break
        
        return chunks


class SentenceChunker(ChunkingStrategy):
    """Sentence-based chunking."""
    
    def __init__(self, max_size: int = 1000):
        config = {'max_size': max_size}
        super().__init__('sentence_based', config)
        self.max_size = max_size
        # Simple sentence boundary detection
        self.sentence_endings = re.compile(r'[.!?]+\s+')
    
    def chunk(self, text: str) -> List[Dict[str, Any]]:
        # Split into sentences
        sentences = self._split_sentences(text)
        chunks = []
        current_chunk = []
        current_size = 0
        chunk_id = 0
        
        for sentence in sentences:
            sentence_size = len(sentence)
            
            if current_size + sentence_size > self.max_size and current_chunk:
                # Save current chunk
                chunk_text = ' '.join(current_chunk)
                chunks.append({
                    'id': chunk_id,
                    'text': chunk_text,
                    'start': 0,  # Approximate
                    'end': len(chunk_text),
                    'size': len(chunk_text),
                    'sentence_count': len(current_chunk)
                })
                chunk_id += 1
                current_chunk = [sentence]
                current_size = sentence_size
            else:
                current_chunk.append(sentence)
                current_size += sentence_size
        
        # Add final chunk
        if current_chunk:
            chunk_text = ' '.join(current_chunk)
            chunks.append({
                'id': chunk_id,
                'text': chunk_text,
                'start': 0,
                'end': len(chunk_text),
                'size': len(chunk_text),
                'sentence_count': len(current_chunk)
            })
        
        return chunks
    
    def _split_sentences(self, text: str) -> List[str]:
        """Simple sentence splitting."""
        sentences = []
        parts = self.sentence_endings.split(text)
        
        for i, part in enumerate(parts[:-1]):
            # Add the sentence ending back
            ending_match = list(self.sentence_endings.finditer(text))
            if i < len(ending_match):
                sentence = part + ending_match[i].group().strip()
            else:
                sentence = part
            
            if sentence.strip():
                sentences.append(sentence.strip())
        
        # Add final part if it exists
        if parts[-1].strip():
            sentences.append(parts[-1].strip())
        
        return [s for s in sentences if len(s.strip()) > 0]


class ParagraphChunker(ChunkingStrategy):
    """Paragraph-based chunking."""
    
    def __init__(self, max_size: int = 2000, min_paragraph_size: int = 50):
        config = {'max_size': max_size, 'min_paragraph_size': min_paragraph_size}
        super().__init__('paragraph_based', config)
        self.max_size = max_size
        self.min_paragraph_size = min_paragraph_size
    
    def chunk(self, text: str) -> List[Dict[str, Any]]:
        # Split by double newlines (paragraph boundaries)
        paragraphs = [p.strip() for p in re.split(r'\n\s*\n', text) if p.strip()]
        chunks = []
        current_chunk = []
        current_size = 0
        chunk_id = 0
        
        for paragraph in paragraphs:
            paragraph_size = len(paragraph)
            
            # Skip very short paragraphs unless they're the only content
            if paragraph_size < self.min_paragraph_size and len(paragraphs) > 1:
                continue
            
            if current_size + paragraph_size > self.max_size and current_chunk:
                # Save current chunk
                chunk_text = '\n\n'.join(current_chunk)
                chunks.append({
                    'id': chunk_id,
                    'text': chunk_text,
                    'start': 0,
                    'end': len(chunk_text),
                    'size': len(chunk_text),
                    'paragraph_count': len(current_chunk)
                })
                chunk_id += 1
                current_chunk = [paragraph]
                current_size = paragraph_size
            else:
                current_chunk.append(paragraph)
                current_size += paragraph_size + 2  # Account for newlines
        
        # Add final chunk
        if current_chunk:
            chunk_text = '\n\n'.join(current_chunk)
            chunks.append({
                'id': chunk_id,
                'text': chunk_text,
                'start': 0,
                'end': len(chunk_text),
                'size': len(chunk_text),
                'paragraph_count': len(current_chunk)
            })
        
        return chunks


class SemanticChunker(ChunkingStrategy):
    """Heading-aware semantic chunking."""
    
    def __init__(self, max_size: int = 1500, heading_weight: float = 2.0):
        config = {'max_size': max_size, 'heading_weight': heading_weight}
        super().__init__('semantic_heading', config)
        self.max_size = max_size
        self.heading_weight = heading_weight
        
        # Markdown and plain text heading patterns
        self.heading_patterns = [
            re.compile(r'^#{1,6}\s+(.+)$', re.MULTILINE),  # Markdown headers
            re.compile(r'^(.+)\n[=-]+\s*$', re.MULTILINE),  # Underlined headers
            re.compile(r'^\d+\.\s*(.+)$', re.MULTILINE),   # Numbered sections
        ]
    
    def chunk(self, text: str) -> List[Dict[str, Any]]:
        sections = self._identify_sections(text)
        chunks = []
        chunk_id = 0
        
        for section in sections:
            section_chunks = self._chunk_section(section, chunk_id)
            chunks.extend(section_chunks)
            chunk_id += len(section_chunks)
        
        return chunks
    
    def _identify_sections(self, text: str) -> List[Dict[str, Any]]:
        """Identify sections based on headings."""
        sections = []
        lines = text.split('\n')
        current_section = {'heading': 'Introduction', 'content': '', 'level': 0}
        
        for line in lines:
            is_heading = False
            heading_level = 0
            heading_text = line.strip()
            
            # Check for markdown headers
            if line.strip().startswith('#'):
                level = len(line) - len(line.lstrip('#'))
                if level <= 6:
                    heading_text = line.strip('#').strip()
                    heading_level = level
                    is_heading = True
            
            # Check for underlined headers  
            elif len(sections) > 0 and line.strip() and all(c in '=-' for c in line.strip()):
                # Previous line might be heading
                if current_section['content']:
                    content_lines = current_section['content'].strip().split('\n')
                    if content_lines:
                        potential_heading = content_lines[-1].strip()
                        if len(potential_heading) > 0 and len(potential_heading) < 100:
                            # Treat as heading
                            current_section['content'] = '\n'.join(content_lines[:-1])
                            sections.append(current_section)
                            current_section = {
                                'heading': potential_heading,
                                'content': '',
                                'level': 1 if '=' in line else 2
                            }
                            continue
            
            if is_heading:
                if current_section['content'].strip():
                    sections.append(current_section)
                current_section = {
                    'heading': heading_text,
                    'content': '',
                    'level': heading_level
                }
            else:
                current_section['content'] += line + '\n'
        
        # Add final section
        if current_section['content'].strip():
            sections.append(current_section)
        
        return sections
    
    def _chunk_section(self, section: Dict[str, Any], start_id: int) -> List[Dict[str, Any]]:
        """Chunk a single section."""
        content = section['content'].strip()
        if not content:
            return []
        
        heading = section['heading']
        chunks = []
        
        # If section is small enough, return as single chunk
        if len(content) <= self.max_size:
            chunks.append({
                'id': start_id,
                'text': f"{heading}\n\n{content}" if heading else content,
                'start': 0,
                'end': len(content),
                'size': len(content),
                'heading': heading,
                'level': section['level']
            })
            return chunks
        
        # Split large sections by paragraphs
        paragraphs = [p.strip() for p in content.split('\n\n') if p.strip()]
        current_chunk = []
        current_size = len(heading) + 2 if heading else 0  # Account for heading
        chunk_id = start_id
        
        for paragraph in paragraphs:
            paragraph_size = len(paragraph)
            
            if current_size + paragraph_size > self.max_size and current_chunk:
                # Save current chunk
                chunk_text = '\n\n'.join(current_chunk)
                if heading and chunk_id == start_id:
                    chunk_text = f"{heading}\n\n{chunk_text}"
                
                chunks.append({
                    'id': chunk_id,
                    'text': chunk_text,
                    'start': 0,
                    'end': len(chunk_text),
                    'size': len(chunk_text),
                    'heading': heading if chunk_id == start_id else f"{heading} (continued)",
                    'level': section['level']
                })
                chunk_id += 1
                current_chunk = [paragraph]
                current_size = paragraph_size
            else:
                current_chunk.append(paragraph)
                current_size += paragraph_size + 2  # Account for newlines
        
        # Add final chunk
        if current_chunk:
            chunk_text = '\n\n'.join(current_chunk)
            if heading and chunk_id == start_id:
                chunk_text = f"{heading}\n\n{chunk_text}"
            elif heading:
                chunk_text = f"{heading} (continued)\n\n{chunk_text}"
            
            chunks.append({
                'id': chunk_id,
                'text': chunk_text,
                'start': 0,
                'end': len(chunk_text),
                'size': len(chunk_text),
                'heading': heading if chunk_id == start_id else f"{heading} (continued)",
                'level': section['level']
            })
        
        return chunks


class ChunkAnalyzer:
    """Analyzes chunks and provides quality metrics."""
    
    def __init__(self):
        self.vocabulary = set()
        self.word_freq = Counter()
    
    def analyze_chunks(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Comprehensive chunk analysis."""
        if not chunks:
            return {'error': 'No chunks to analyze'}
        
        sizes = [chunk['size'] for chunk in chunks]
        
        # Basic size statistics
        size_stats = {
            'count': len(chunks),
            'mean': statistics.mean(sizes),
            'median': statistics.median(sizes),
            'std': statistics.stdev(sizes) if len(sizes) > 1 else 0,
            'min': min(sizes),
            'max': max(sizes),
            'total': sum(sizes)
        }
        
        # Boundary quality analysis
        boundary_quality = self._analyze_boundary_quality(chunks)
        
        # Semantic coherence (simple heuristic)
        coherence_score = self._calculate_semantic_coherence(chunks)
        
        # Vocabulary distribution
        vocab_stats = self._analyze_vocabulary(chunks)
        
        return {
            'size_statistics': size_stats,
            'boundary_quality': boundary_quality,
            'semantic_coherence': coherence_score,
            'vocabulary_statistics': vocab_stats
        }
    
    def _analyze_boundary_quality(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Analyze how well chunks respect natural boundaries."""
        sentence_breaks = 0
        word_breaks = 0
        total_chunks = len(chunks)
        
        sentence_endings = re.compile(r'[.!?]\s*$')
        
        for chunk in chunks:
            text = chunk['text'].strip()
            if not text:
                continue
            
            # Check if chunk ends with sentence boundary
            if sentence_endings.search(text):
                sentence_breaks += 1
            
            # Check if chunk ends with word boundary
            if text[-1].isalnum() or text[-1] in '.!?':
                word_breaks += 1
        
        return {
            'sentence_boundary_ratio': sentence_breaks / total_chunks if total_chunks > 0 else 0,
            'word_boundary_ratio': word_breaks / total_chunks if total_chunks > 0 else 0,
            'clean_breaks': sentence_breaks,
            'total_chunks': total_chunks
        }
    
    def _calculate_semantic_coherence(self, chunks: List[Dict[str, Any]]) -> float:
        """Simple semantic coherence heuristic based on vocabulary overlap."""
        if len(chunks) < 2:
            return 1.0
        
        coherence_scores = []
        
        for i in range(len(chunks) - 1):
            chunk1_words = set(re.findall(r'\b\w+\b', chunks[i]['text'].lower()))
            chunk2_words = set(re.findall(r'\b\w+\b', chunks[i+1]['text'].lower()))
            
            if not chunk1_words or not chunk2_words:
                continue
            
            # Jaccard similarity as coherence measure
            intersection = len(chunk1_words & chunk2_words)
            union = len(chunk1_words | chunk2_words)
            
            if union > 0:
                coherence_scores.append(intersection / union)
        
        return statistics.mean(coherence_scores) if coherence_scores else 0.0
    
    def _analyze_vocabulary(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Analyze vocabulary distribution across chunks."""
        all_words = []
        chunk_vocab_sizes = []
        
        for chunk in chunks:
            words = re.findall(r'\b\w+\b', chunk['text'].lower())
            all_words.extend(words)
            chunk_vocab_sizes.append(len(set(words)))
        
        total_vocab = len(set(all_words))
        word_freq = Counter(all_words)
        
        return {
            'total_vocabulary': total_vocab,
            'avg_chunk_vocabulary': statistics.mean(chunk_vocab_sizes) if chunk_vocab_sizes else 0,
            'vocabulary_diversity': total_vocab / len(all_words) if all_words else 0,
            'most_common_words': word_freq.most_common(10)
        }


class ChunkingOptimizer:
    """Main optimizer that tests different chunking strategies."""
    
    def __init__(self):
        self.analyzer = ChunkAnalyzer()
    
    def optimize(self, corpus: DocumentCorpus, config: Dict[str, Any] = None) -> Dict[str, Any]:
        """Test all chunking strategies and recommend the best one."""
        config = config or {}
        
        strategies = self._create_strategies(config)
        results = {}
        
        print(f"Testing {len(strategies)} chunking strategies...")
        
        for strategy in strategies:
            print(f"  Testing {strategy.name}...")
            strategy_results = self._test_strategy(corpus, strategy)
            results[strategy.name] = strategy_results
        
        # Recommend best strategy
        recommendation = self._recommend_strategy(results)
        
        return {
            'corpus_info': {
                'document_count': len(corpus.documents),
                'total_size': sum(d['size'] for d in corpus.documents),
                'avg_document_size': statistics.mean([d['size'] for d in corpus.documents])
            },
            'strategy_results': results,
            'recommendation': recommendation,
            'sample_chunks': self._generate_sample_chunks(corpus, recommendation['best_strategy'])
        }
    
    def _create_strategies(self, config: Dict[str, Any]) -> List[ChunkingStrategy]:
        """Create all chunking strategies to test."""
        strategies = []
        
        # Fixed-size strategies
        for size in config.get('fixed_sizes', [512, 1000, 1500]):
            for overlap in config.get('overlaps', [50, 100]):
                strategies.append(FixedSizeChunker(size, overlap, 'char'))
        
        # Sentence-based strategies
        for max_size in config.get('sentence_max_sizes', [800, 1200]):
            strategies.append(SentenceChunker(max_size))
        
        # Paragraph-based strategies
        for max_size in config.get('paragraph_max_sizes', [1500, 2000]):
            strategies.append(ParagraphChunker(max_size))
        
        # Semantic strategies
        for max_size in config.get('semantic_max_sizes', [1200, 1800]):
            strategies.append(SemanticChunker(max_size))
        
        return strategies
    
    def _test_strategy(self, corpus: DocumentCorpus, strategy: ChunkingStrategy) -> Dict[str, Any]:
        """Test a single chunking strategy."""
        all_chunks = []
        document_results = []
        
        for doc in corpus.documents:
            try:
                chunks = strategy.chunk(doc['content'])
                all_chunks.extend(chunks)
                
                doc_analysis = self.analyzer.analyze_chunks(chunks)
                document_results.append({
                    'path': doc['path'],
                    'chunk_count': len(chunks),
                    'analysis': doc_analysis
                })
            except Exception as e:
                print(f"    Error processing {doc['path']}: {e}")
                continue
        
        # Overall analysis
        overall_analysis = self.analyzer.analyze_chunks(all_chunks)
        
        return {
            'strategy_config': strategy.config,
            'total_chunks': len(all_chunks),
            'overall_analysis': overall_analysis,
            'document_results': document_results,
            'performance_score': self._calculate_performance_score(overall_analysis)
        }
    
    def _calculate_performance_score(self, analysis: Dict[str, Any]) -> float:
        """Calculate overall performance score for a strategy."""
        if 'error' in analysis:
            return 0.0
        
        size_stats = analysis['size_statistics']
        boundary_quality = analysis['boundary_quality']
        coherence = analysis['semantic_coherence']
        
        # Normalize metrics to 0-1 range and combine
        size_consistency = 1.0 - min(size_stats['std'] / size_stats['mean'], 1.0) if size_stats['mean'] > 0 else 0
        boundary_score = (boundary_quality['sentence_boundary_ratio'] + boundary_quality['word_boundary_ratio']) / 2
        coherence_score = coherence
        
        # Weighted combination
        return (size_consistency * 0.3 + boundary_score * 0.4 + coherence_score * 0.3)
    
    def _recommend_strategy(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """Recommend the best chunking strategy based on analysis."""
        best_strategy = None
        best_score = 0
        
        strategy_scores = {}
        
        for strategy_name, result in results.items():
            score = result['performance_score']
            strategy_scores[strategy_name] = score
            
            if score > best_score:
                best_score = score
                best_strategy = strategy_name
        
        return {
            'best_strategy': best_strategy,
            'best_score': best_score,
            'all_scores': strategy_scores,
            'reasoning': self._generate_reasoning(best_strategy, results[best_strategy] if best_strategy else None)
        }
    
    def _generate_reasoning(self, strategy_name: str, result: Dict[str, Any]) -> str:
        """Generate human-readable reasoning for the recommendation."""
        if not result:
            return "No valid strategy found."
        
        analysis = result['overall_analysis']
        size_stats = analysis['size_statistics']
        boundary = analysis['boundary_quality']
        
        reasoning = f"Recommended '{strategy_name}' because:\n"
        reasoning += f"- Average chunk size: {size_stats['mean']:.0f} characters\n"
        reasoning += f"- Size consistency: {size_stats['std']:.0f} std deviation\n"
        reasoning += f"- Boundary quality: {boundary['sentence_boundary_ratio']:.2%} clean sentence breaks\n"
        reasoning += f"- Semantic coherence: {analysis['semantic_coherence']:.3f}\n"
        
        return reasoning
    
    def _generate_sample_chunks(self, corpus: DocumentCorpus, strategy_name: str) -> List[Dict[str, Any]]:
        """Generate sample chunks using the recommended strategy."""
        if not strategy_name or not corpus.documents:
            return []
        
        # Create strategy instance
        strategy = None
        if 'fixed_size' in strategy_name:
            strategy = FixedSizeChunker()
        elif 'sentence' in strategy_name:
            strategy = SentenceChunker()
        elif 'paragraph' in strategy_name:
            strategy = ParagraphChunker()
        elif 'semantic' in strategy_name:
            strategy = SemanticChunker()
        
        if not strategy:
            return []
        
        # Get chunks from first document
        sample_doc = corpus.documents[0]
        chunks = strategy.chunk(sample_doc['content'])
        
        # Return first 3 chunks as samples
        return chunks[:3]


def main():
    """Main function with command-line interface."""
    parser = argparse.ArgumentParser(description='Analyze documents and recommend optimal chunking strategy')
    parser.add_argument('directory', help='Directory containing text/markdown documents')
    parser.add_argument('--output', '-o', help='Output file for results (JSON format)')
    parser.add_argument('--config', '-c', help='Configuration file (JSON format)')
    parser.add_argument('--extensions', nargs='+', default=['.txt', '.md', '.markdown'], 
                       help='File extensions to process')
    parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
    
    args = parser.parse_args()
    
    # Load configuration
    config = {}
    if args.config and os.path.exists(args.config):
        with open(args.config, 'r') as f:
            config = json.load(f)
    
    try:
        # Load corpus
        print(f"Loading documents from {args.directory}...")
        corpus = DocumentCorpus(args.directory, args.extensions)
        
        # Run optimization
        optimizer = ChunkingOptimizer()
        results = optimizer.optimize(corpus, config)
        
        # Save results
        if args.output:
            with open(args.output, 'w') as f:
                json.dump(results, f, indent=2)
            print(f"Results saved to {args.output}")
        
        # Print summary
        print("\n" + "="*60)
        print("CHUNKING OPTIMIZATION RESULTS")
        print("="*60)
        
        corpus_info = results['corpus_info']
        print(f"Corpus: {corpus_info['document_count']} documents, {corpus_info['total_size']:,} characters")
        
        recommendation = results['recommendation']
        print(f"\nRecommended Strategy: {recommendation['best_strategy']}")
        print(f"Performance Score: {recommendation['best_score']:.3f}")
        print(f"\nReasoning:\n{recommendation['reasoning']}")
        
        if args.verbose:
            print("\nAll Strategy Scores:")
            for strategy, score in recommendation['all_scores'].items():
                print(f"  {strategy}: {score:.3f}")
        
        print("\nSample Chunks:")
        for i, chunk in enumerate(results['sample_chunks'][:2]):
            print(f"\nChunk {i+1} ({chunk['size']} chars):")
            print("-" * 40)
            print(chunk['text'][:200] + "..." if len(chunk['text']) > 200 else chunk['text'])
        
    except Exception as e:
        print(f"Error: {e}")
        return 1
    
    return 0


if __name__ == '__main__':
    exit(main())

#!/usr/bin/env python3
"""
RAG Pipeline Designer - Designs complete RAG pipelines based on requirements.

This script analyzes requirements and generates a comprehensive RAG pipeline design
including architecture diagrams, component recommendations, configuration templates,
and cost projections.

Components designed:
- Chunking strategy recommendation
- Embedding model selection
- Vector database recommendation  
- Retrieval approach (dense/sparse/hybrid)
- Reranking configuration
- Evaluation framework setup
- Production deployment patterns

No external dependencies - uses only Python standard library.
"""

import argparse
import json
import math
import os
from typing import Dict, List, Tuple, Any, Optional
from dataclasses import dataclass, asdict
from enum import Enum


class Scale(Enum):
    """System scale categories."""
    SMALL = "small"      # < 1M documents, < 1K queries/day
    MEDIUM = "medium"    # 1M-100M documents, 1K-100K queries/day  
    LARGE = "large"      # 100M+ documents, 100K+ queries/day


class DocumentType(Enum):
    """Document type categories."""
    TEXT = "text"                # Plain text, articles
    TECHNICAL = "technical"      # Documentation, manuals
    CODE = "code"               # Source code files
    SCIENTIFIC = "scientific"    # Research papers, journals
    LEGAL = "legal"             # Legal documents, contracts
    MIXED = "mixed"             # Multiple document types


class Latency(Enum):
    """Latency requirements."""
    REAL_TIME = "real_time"      # < 100ms
    INTERACTIVE = "interactive"   # < 500ms  
    BATCH = "batch"              # > 1s acceptable


@dataclass
class Requirements:
    """RAG system requirements."""
    document_types: List[str]
    document_count: int
    avg_document_size: int  # characters
    queries_per_day: int
    query_patterns: List[str]  # e.g., ["factual", "conversational", "analytical"]
    latency_requirement: str
    budget_monthly: float  # USD
    accuracy_priority: float  # 0-1 scale
    cost_priority: float     # 0-1 scale
    maintenance_complexity: str  # "low", "medium", "high"


@dataclass
class ComponentRecommendation:
    """Recommendation for a pipeline component."""
    name: str
    type: str
    config: Dict[str, Any]
    rationale: str
    pros: List[str]
    cons: List[str]
    cost_monthly: float


@dataclass
class PipelineDesign:
    """Complete RAG pipeline design."""
    chunking: ComponentRecommendation
    embedding: ComponentRecommendation  
    vector_db: ComponentRecommendation
    retrieval: ComponentRecommendation
    reranking: Optional[ComponentRecommendation]
    evaluation: ComponentRecommendation
    total_cost: float
    architecture_diagram: str
    config_templates: Dict[str, Any]


class RAGPipelineDesigner:
    """Main pipeline designer class."""
    
    def __init__(self):
        self.embedding_models = self._load_embedding_models()
        self.vector_databases = self._load_vector_databases()
        self.chunking_strategies = self._load_chunking_strategies()
    
    def design_pipeline(self, requirements: Requirements) -> PipelineDesign:
        """Design complete RAG pipeline based on requirements."""
        print(f"Designing RAG pipeline for {requirements.document_count:,} documents...")
        
        # Determine system scale
        scale = self._determine_scale(requirements)
        print(f"System scale: {scale.value}")
        
        # Design each component
        chunking = self._recommend_chunking(requirements, scale)
        embedding = self._recommend_embedding(requirements, scale)  
        vector_db = self._recommend_vector_db(requirements, scale)
        retrieval = self._recommend_retrieval(requirements, scale)
        reranking = self._recommend_reranking(requirements, scale)
        evaluation = self._recommend_evaluation(requirements, scale)
        
        # Calculate total cost
        total_cost = (chunking.cost_monthly + embedding.cost_monthly + 
                     vector_db.cost_monthly + retrieval.cost_monthly + 
                     evaluation.cost_monthly)
        if reranking:
            total_cost += reranking.cost_monthly
        
        # Generate architecture diagram
        architecture = self._generate_architecture_diagram(
            chunking, embedding, vector_db, retrieval, reranking, evaluation
        )
        
        # Generate configuration templates
        configs = self._generate_config_templates(
            chunking, embedding, vector_db, retrieval, reranking, evaluation
        )
        
        return PipelineDesign(
            chunking=chunking,
            embedding=embedding,
            vector_db=vector_db,
            retrieval=retrieval,
            reranking=reranking,
            evaluation=evaluation,
            total_cost=total_cost,
            architecture_diagram=architecture,
            config_templates=configs
        )
    
    def _determine_scale(self, req: Requirements) -> Scale:
        """Determine system scale based on requirements."""
        if req.document_count < 1_000_000 and req.queries_per_day < 1_000:
            return Scale.SMALL
        elif req.document_count < 100_000_000 and req.queries_per_day < 100_000:
            return Scale.MEDIUM
        else:
            return Scale.LARGE
    
    def _recommend_chunking(self, req: Requirements, scale: Scale) -> ComponentRecommendation:
        """Recommend chunking strategy."""
        doc_types = set(req.document_types)
        
        if "code" in doc_types:
            strategy = "semantic_code_aware"
            config = {"max_size": 1000, "preserve_functions": True, "overlap": 50}
            rationale = "Code documents benefit from function/class boundary awareness"
        elif "technical" in doc_types or "scientific" in doc_types:
            strategy = "semantic_heading_aware" 
            config = {"max_size": 1500, "heading_weight": 2.0, "overlap": 100}
            rationale = "Technical documents have clear hierarchical structure"
        elif len(doc_types) > 2 or "mixed" in doc_types:
            strategy = "adaptive_chunking"
            config = {"strategies": ["paragraph", "sentence", "fixed"], "auto_select": True}
            rationale = "Mixed document types require adaptive strategy selection"
        else:
            if req.avg_document_size > 5000:
                strategy = "paragraph_based"
                config = {"max_size": 2000, "min_paragraph_size": 100}
                rationale = "Large documents benefit from paragraph-based chunking"
            else:
                strategy = "sentence_based"  
                config = {"max_size": 1000, "sentence_overlap": 1}
                rationale = "Small to medium documents work well with sentence chunking"
        
        return ComponentRecommendation(
            name=strategy,
            type="chunking",
            config=config,
            rationale=rationale,
            pros=self._get_chunking_pros(strategy),
            cons=self._get_chunking_cons(strategy),
            cost_monthly=0.0  # Processing cost only
        )
    
    def _recommend_embedding(self, req: Requirements, scale: Scale) -> ComponentRecommendation:
        """Recommend embedding model."""
        doc_types = set(req.document_types)
        
        # Consider accuracy vs cost priority
        high_accuracy = req.accuracy_priority > 0.7
        cost_sensitive = req.cost_priority > 0.6
        
        if "code" in doc_types:
            if high_accuracy and not cost_sensitive:
                model = "openai-code-search-ada-002"
                cost_per_1k_tokens = 0.0001
                dimensions = 1536
            else:
                model = "sentence-transformers/code-bert-base"
                cost_per_1k_tokens = 0.0  # Self-hosted
                dimensions = 768
        elif "scientific" in doc_types:
            if high_accuracy:
                model = "openai-text-embedding-ada-002"
                cost_per_1k_tokens = 0.0001
                dimensions = 1536
            else:
                model = "sentence-transformers/scibert-nli" 
                cost_per_1k_tokens = 0.0
                dimensions = 768
        else:
            if cost_sensitive or scale == Scale.SMALL:
                model = "sentence-transformers/all-MiniLM-L6-v2"
                cost_per_1k_tokens = 0.0
                dimensions = 384
            elif high_accuracy:
                model = "openai-text-embedding-ada-002"
                cost_per_1k_tokens = 0.0001  
                dimensions = 1536
            else:
                model = "sentence-transformers/all-mpnet-base-v2"
                cost_per_1k_tokens = 0.0
                dimensions = 768
        
        # Calculate monthly embedding cost
        total_tokens = req.document_count * (req.avg_document_size / 4)  # ~4 chars per token
        query_tokens = req.queries_per_day * 30 * 20  # ~20 tokens per query per month
        monthly_cost = (total_tokens + query_tokens) * cost_per_1k_tokens / 1000
        
        return ComponentRecommendation(
            name=model,
            type="embedding",
            config={
                "model": model,
                "dimensions": dimensions,
                "batch_size": 100 if scale == Scale.SMALL else 1000,
                "cache_embeddings": True
            },
            rationale=f"Selected for {doc_types} with accuracy priority {req.accuracy_priority}",
            pros=self._get_embedding_pros(model),
            cons=self._get_embedding_cons(model),
            cost_monthly=monthly_cost
        )
    
    def _recommend_vector_db(self, req: Requirements, scale: Scale) -> ComponentRecommendation:
        """Recommend vector database."""
        if scale == Scale.SMALL and req.cost_priority > 0.7:
            db = "chroma"
            cost = 0.0
            rationale = "Local/embedded database suitable for small scale and cost optimization"
        elif scale == Scale.SMALL and req.maintenance_complexity == "low":
            db = "pgvector"
            cost = 50.0  # PostgreSQL hosting
            rationale = "Leverage existing PostgreSQL infrastructure"
        elif scale == Scale.LARGE or req.latency_requirement == "real_time":
            db = "pinecone"
            vectors = req.document_count * 2  # Account for chunking
            cost = max(70, vectors * 0.00005)  # $70 base + $0.00005 per vector
            rationale = "Managed service with excellent performance for large scale"
        elif req.maintenance_complexity == "low":
            db = "weaviate_cloud"
            vectors = req.document_count * 2
            cost = max(25, vectors * 0.00003)
            rationale = "Managed Weaviate with good balance of features and cost"
        else:
            db = "qdrant"
            cost = 100.0  # Self-hosted infrastructure estimate
            rationale = "High performance self-hosted option with good scaling"
        
        return ComponentRecommendation(
            name=db,
            type="vector_database",
            config=self._get_vector_db_config(db, req, scale),
            rationale=rationale,
            pros=self._get_vector_db_pros(db),
            cons=self._get_vector_db_cons(db),
            cost_monthly=cost
        )
    
    def _recommend_retrieval(self, req: Requirements, scale: Scale) -> ComponentRecommendation:
        """Recommend retrieval strategy."""
        if req.accuracy_priority > 0.8:
            strategy = "hybrid"
            rationale = "Hybrid retrieval for maximum accuracy combining dense and sparse methods"
        elif "technical" in req.document_types or "code" in req.document_types:
            strategy = "hybrid"
            rationale = "Technical content benefits from both semantic and keyword matching"
        elif req.latency_requirement == "real_time":
            strategy = "dense"
            rationale = "Dense retrieval faster for real-time requirements"
        else:
            strategy = "dense"
            rationale = "Dense retrieval suitable for general text search"
        
        return ComponentRecommendation(
            name=strategy,
            type="retrieval", 
            config={
                "strategy": strategy,
                "dense_weight": 0.7 if strategy == "hybrid" else 1.0,
                "sparse_weight": 0.3 if strategy == "hybrid" else 0.0,
                "top_k": 20 if req.accuracy_priority > 0.7 else 10,
                "similarity_threshold": 0.7
            },
            rationale=rationale,
            pros=self._get_retrieval_pros(strategy),
            cons=self._get_retrieval_cons(strategy),
            cost_monthly=0.0
        )
    
    def _recommend_reranking(self, req: Requirements, scale: Scale) -> Optional[ComponentRecommendation]:
        """Recommend reranking if beneficial."""
        if req.accuracy_priority < 0.6 or req.latency_requirement == "real_time":
            return None
        
        if req.cost_priority > 0.8:
            return None
        
        # Estimate reranking queries per month
        monthly_queries = req.queries_per_day * 30
        cost_per_query = 0.002  # Estimated cost for cross-encoder reranking
        monthly_cost = monthly_queries * cost_per_query
        
        if monthly_cost > req.budget_monthly * 0.3:  # Don't exceed 30% of budget
            return None
        
        return ComponentRecommendation(
            name="cross_encoder_reranking",
            type="reranking",
            config={
                "model": "cross-encoder/ms-marco-MiniLM-L-12-v2",
                "rerank_top_k": 20,
                "return_top_k": 5,
                "batch_size": 16
            },
            rationale="Reranking improves precision for high-accuracy requirements",
            pros=["Higher precision", "Better ranking quality", "Handles complex queries"],
            cons=["Additional latency", "Higher cost", "More complexity"],
            cost_monthly=monthly_cost
        )
    
    def _recommend_evaluation(self, req: Requirements, scale: Scale) -> ComponentRecommendation:
        """Recommend evaluation framework."""
        return ComponentRecommendation(
            name="comprehensive_evaluation",
            type="evaluation",
            config={
                "metrics": ["precision@k", "recall@k", "mrr", "ndcg"],
                "k_values": [1, 3, 5, 10],
                "faithfulness_check": True,
                "relevance_scoring": True,
                "evaluation_frequency": "weekly" if scale == Scale.LARGE else "monthly",
                "sample_size": min(1000, req.queries_per_day * 7)
            },
            rationale="Comprehensive evaluation essential for production RAG systems",
            pros=["Quality monitoring", "Performance tracking", "Issue detection"],
            cons=["Additional overhead", "Requires ground truth data"],
            cost_monthly=20.0  # Evaluation tooling and compute
        )
    
    def _generate_architecture_diagram(self, chunking: ComponentRecommendation, 
                                     embedding: ComponentRecommendation,
                                     vector_db: ComponentRecommendation,
                                     retrieval: ComponentRecommendation,
                                     reranking: Optional[ComponentRecommendation],
                                     evaluation: ComponentRecommendation) -> str:
        """Generate Mermaid architecture diagram."""
        
        diagram = """```mermaid
graph TB
    %% Document Processing Pipeline
    A[Document Corpus] --> B[Document Chunking]
    B --> C[Embedding Generation]
    C --> D[Vector Database Storage]
    
    %% Query Processing Pipeline  
    E[User Query] --> F[Query Processing]
    F --> G[Vector Search]
    D --> G
    G --> H[Retrieved Chunks]
"""
        
        if reranking:
            diagram += "    H --> I[Reranking]\n    I --> J[Final Results]\n"
        else:
            diagram += "    H --> J[Final Results]\n"
        
        diagram += """    
    %% Evaluation Pipeline
    J --> K[Response Generation]
    K --> L[Evaluation Metrics]
    
    %% Component Details
    B -.-> B1[Strategy: """ + chunking.name + """]
    C -.-> C1[Model: """ + embedding.name + """]
    D -.-> D1[Database: """ + vector_db.name + """]
    G -.-> G1[Method: """ + retrieval.name + """]
"""
        
        if reranking:
            diagram += "    I -.-> I1[Model: " + reranking.name + "]\n"
        
        diagram += "    L -.-> L1[Framework: " + evaluation.name + "]\n```"
        
        return diagram
    
    def _generate_config_templates(self, *components) -> Dict[str, Any]:
        """Generate configuration templates for all components."""
        configs = {}
        
        for component in components:
            if component:
                configs[component.type] = {
                    "component": component.name,
                    "config": component.config,
                    "rationale": component.rationale
                }
        
        # Add deployment configuration
        configs["deployment"] = {
            "infrastructure": "cloud" if any("pinecone" in str(c.name) for c in components if c) else "hybrid",
            "scaling": {
                "auto_scaling": True,
                "min_replicas": 1,
                "max_replicas": 10
            },
            "monitoring": {
                "metrics": ["latency", "throughput", "accuracy"],
                "alerts": ["high_latency", "low_accuracy", "service_down"]
            }
        }
        
        return configs
    
    def _load_embedding_models(self) -> Dict[str, Dict[str, Any]]:
        """Load embedding model specifications."""
        return {
            "openai-text-embedding-ada-002": {
                "dimensions": 1536,
                "cost_per_1k_tokens": 0.0001,
                "quality": "high",
                "speed": "medium"
            },
            "sentence-transformers/all-mpnet-base-v2": {
                "dimensions": 768, 
                "cost_per_1k_tokens": 0.0,
                "quality": "high",
                "speed": "medium"
            },
            "sentence-transformers/all-MiniLM-L6-v2": {
                "dimensions": 384,
                "cost_per_1k_tokens": 0.0,
                "quality": "medium",
                "speed": "fast"
            }
        }
    
    def _load_vector_databases(self) -> Dict[str, Dict[str, Any]]:
        """Load vector database specifications."""
        return {
            "pinecone": {"managed": True, "scaling": "excellent", "cost": "high"},
            "weaviate": {"managed": False, "scaling": "good", "cost": "medium"},
            "qdrant": {"managed": False, "scaling": "excellent", "cost": "low"},
            "chroma": {"managed": False, "scaling": "poor", "cost": "free"},
            "pgvector": {"managed": False, "scaling": "good", "cost": "medium"}
        }
    
    def _load_chunking_strategies(self) -> Dict[str, Dict[str, Any]]:
        """Load chunking strategy specifications."""
        return {
            "fixed_size": {"complexity": "low", "quality": "medium"},
            "sentence_based": {"complexity": "medium", "quality": "good"}, 
            "paragraph_based": {"complexity": "medium", "quality": "good"},
            "semantic_heading_aware": {"complexity": "high", "quality": "excellent"}
        }
    
    def _get_vector_db_config(self, db: str, req: Requirements, scale: Scale) -> Dict[str, Any]:
        """Get vector database configuration."""
        base_config = {
            "collection_name": "rag_documents",
            "distance_metric": "cosine",
            "index_type": "hnsw"
        }
        
        if db == "pinecone":
            base_config.update({
                "environment": "us-east1-gcp",
                "replicas": 1 if scale == Scale.SMALL else 2,
                "shards": 1 if scale != Scale.LARGE else 3
            })
        elif db == "qdrant":
            base_config.update({
                "memory_mapping": True,
                "quantization": scale == Scale.LARGE,
                "replication_factor": 1 if scale == Scale.SMALL else 2
            })
        
        return base_config
    
    def _get_chunking_pros(self, strategy: str) -> List[str]:
        """Get pros for chunking strategy."""
        pros_map = {
            "semantic_heading_aware": ["Preserves document structure", "High semantic coherence", "Good for technical docs"],
            "paragraph_based": ["Respects natural boundaries", "Good balance", "Readable chunks"],
            "sentence_based": ["Natural language boundaries", "Consistent quality", "Good for general text"],
            "fixed_size": ["Predictable sizes", "Simple implementation", "Consistent processing"],
            "adaptive_chunking": ["Handles mixed content", "Optimizes per document", "Best quality"]
        }
        return pros_map.get(strategy, ["Good general purpose strategy"])
    
    def _get_chunking_cons(self, strategy: str) -> List[str]:
        """Get cons for chunking strategy."""
        cons_map = {
            "semantic_heading_aware": ["Complex implementation", "May create large chunks", "Document-dependent"],
            "paragraph_based": ["Variable sizes", "May break context", "Document-dependent"],
            "sentence_based": ["May create small chunks", "Sentence detection issues", "Variable sizes"],
            "fixed_size": ["Breaks semantic boundaries", "May split sentences", "Context loss"],
            "adaptive_chunking": ["High complexity", "Slower processing", "Harder to debug"]
        }
        return cons_map.get(strategy, ["May not fit all use cases"])
    
    def _get_embedding_pros(self, model: str) -> List[str]:
        """Get pros for embedding model."""
        if "openai" in model:
            return ["High quality", "Regular updates", "Good performance"]
        elif "all-mpnet" in model:
            return ["High quality", "Free to use", "Good balance"]
        elif "MiniLM" in model:
            return ["Fast processing", "Small size", "Good for real-time"]
        else:
            return ["Specialized for domain", "Good performance"]
    
    def _get_embedding_cons(self, model: str) -> List[str]:
        """Get cons for embedding model."""
        if "openai" in model:
            return ["API costs", "Vendor lock-in", "Rate limits"]
        elif "sentence-transformers" in model:
            return ["Self-hosting required", "Model updates needed", "GPU beneficial"]
        else:
            return ["May require fine-tuning", "Domain-specific"]
    
    def _get_vector_db_pros(self, db: str) -> List[str]:
        """Get pros for vector database."""
        pros_map = {
            "pinecone": ["Fully managed", "Excellent performance", "Auto-scaling"],
            "weaviate": ["Rich features", "GraphQL API", "Multi-modal"],
            "qdrant": ["High performance", "Rust-based", "Good scaling"],
            "chroma": ["Simple setup", "Free", "Good for development"],
            "pgvector": ["SQL integration", "ACID compliance", "Familiar"]
        }
        return pros_map.get(db, ["Good performance"])
    
    def _get_vector_db_cons(self, db: str) -> List[str]:
        """Get cons for vector database."""
        cons_map = {
            "pinecone": ["Expensive", "Vendor lock-in", "Limited customization"],
            "weaviate": ["Complex setup", "Learning curve", "Resource intensive"],
            "qdrant": ["Self-managed", "Smaller community", "Setup complexity"],
            "chroma": ["Limited scaling", "Not production-ready", "Basic features"],
            "pgvector": ["PostgreSQL knowledge needed", "Less specialized", "Manual optimization"]
        }
        return cons_map.get(db, ["Requires maintenance"])
    
    def _get_retrieval_pros(self, strategy: str) -> List[str]:
        """Get pros for retrieval strategy."""
        pros_map = {
            "dense": ["Semantic understanding", "Good for paraphrases", "Fast"],
            "sparse": ["Exact matching", "Interpretable", "Good for keywords"], 
            "hybrid": ["Best of both", "High accuracy", "Robust"]
        }
        return pros_map.get(strategy, ["Good performance"])
    
    def _get_retrieval_cons(self, strategy: str) -> List[str]:
        """Get cons for retrieval strategy."""
        cons_map = {
            "dense": ["May miss exact matches", "Embedding dependent", "Less interpretable"],
            "sparse": ["Vocabulary mismatch", "No semantic understanding", "Synonym issues"],
            "hybrid": ["More complex", "Tuning required", "Higher latency"]
        }
        return cons_map.get(strategy, ["May require tuning"])


def load_requirements(file_path: str) -> Requirements:
    """Load requirements from JSON file."""
    with open(file_path, 'r') as f:
        data = json.load(f)
    
    return Requirements(**data)


def save_design(design: PipelineDesign, output_path: str):
    """Save pipeline design to JSON file."""
    # Convert to dict for JSON serialization
    design_dict = {}
    
    for field_name in design.__dataclass_fields__:
        value = getattr(design, field_name)
        if isinstance(value, ComponentRecommendation):
            design_dict[field_name] = asdict(value)
        elif value is None:
            design_dict[field_name] = None
        else:
            design_dict[field_name] = value
    
    with open(output_path, 'w') as f:
        json.dump(design_dict, f, indent=2)


def print_design_summary(design: PipelineDesign):
    """Print human-readable design summary."""
    print("\n" + "="*60)
    print("RAG PIPELINE DESIGN SUMMARY")
    print("="*60)
    
    print(f"\n💰 Total Monthly Cost: ${design.total_cost:.2f}")
    
    print(f"\n🔧 Component Recommendations:")
    components = [design.chunking, design.embedding, design.vector_db, 
                 design.retrieval, design.reranking, design.evaluation]
    
    for component in components:
        if component:
            print(f"\n  {component.type.upper()}: {component.name}")
            print(f"    Rationale: {component.rationale}")
            if component.cost_monthly > 0:
                print(f"    Monthly Cost: ${component.cost_monthly:.2f}")
    
    print(f"\n📊 Architecture Diagram:")
    print(design.architecture_diagram)


def main():
    """Main function with command-line interface."""
    parser = argparse.ArgumentParser(description='Design RAG pipeline based on requirements')
    parser.add_argument('requirements', help='JSON file containing system requirements')
    parser.add_argument('--output', '-o', help='Output file for pipeline design (JSON)')
    parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
    
    args = parser.parse_args()
    
    try:
        # Load requirements
        print("Loading requirements...")
        requirements = load_requirements(args.requirements)
        
        # Design pipeline
        designer = RAGPipelineDesigner()
        design = designer.design_pipeline(requirements)
        
        # Save design
        if args.output:
            save_design(design, args.output)
            print(f"Pipeline design saved to {args.output}")
        
        # Print summary
        print_design_summary(design)
        
        if args.verbose:
            print(f"\n📋 Configuration Templates:")
            for component_type, config in design.config_templates.items():
                print(f"\n  {component_type.upper()}:")
                print(f"    {json.dumps(config, indent=4)}")
        
    except Exception as e:
        print(f"Error: {e}")
        return 1
    
    return 0


if __name__ == '__main__':
    exit(main())