"""Core documentation generation logic."""

import logging
from pathlib import Path
from typing import List, Optional, Set
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn

from .config import Config
from .file_processor import FileProcessor, FileInfo
from .openai_client import OpenAIClient

logger = logging.getLogger(__name__)
console = Console()


class DocumentationGenerator:
    """Main class for generating documentation using OpenAI API."""
    
    def __init__(self, config: Config, api_key: Optional[str] = None):
        """Initialize documentation generator."""
        self.config = config
        self.file_processor = FileProcessor(config)
        self.openai_client = OpenAIClient(config, api_key)
    
    def generate_documentation(self, repository_path: Path, output_path: Optional[Path] = None, 
                             additional_skip_patterns: Set[str] = None) -> str:
        """Generate documentation for a repository."""
        console.print(f"[bold blue]Generating documentation for: {repository_path}[/bold blue]")
        
        # Get system prompt
        system_prompt = self.config.get_system_prompt()
        
        # Discover and process files
        console.print("[yellow]Discovering files...[/yellow]")
        files = self.file_processor.discover_files(repository_path, additional_skip_patterns)
        
        if not files:
            console.print("[red]No files found to process![/red]")
            return "No files found to generate documentation for."
        
        # Display files that will be processed
        console.print(f"[cyan]Found {len(files)} files to process:[/cyan]")
        for file_info in files:
            console.print(f"  • {file_info.relative_path} ({file_info.size} bytes)")
        
        # Read file contents
        console.print("[yellow]Reading file contents...[/yellow]")
        files = self.file_processor.read_all_files(files)
        
        # Estimate tokens and cost
        total_content = self.file_processor.format_files_for_documentation(files)
        estimated_tokens = self.openai_client.count_tokens(total_content)
        estimated_cost = self.openai_client.estimate_cost(estimated_tokens)
        
        console.print(f"[cyan]Estimated tokens: {estimated_tokens:,}[/cyan]")
        console.print(f"[cyan]Estimated cost: ${estimated_cost:.4f}[/cyan]")
        
        # Generate documentation
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Generating...", total=None)
            
            if estimated_tokens <= self.config.max_tokens_per_chunk:
                documentation = self._generate_single_chunk(total_content, system_prompt)
            else:
                documentation = self._generate_multiple_chunks(files, system_prompt)
            
            progress.update(task, completed=True)
        
        # Save documentation
        if output_path:
            self._save_documentation(documentation, output_path)
        
        console.print("[green]Documentation generated successfully![/green]")
        return documentation
    
    def _generate_single_chunk(self, content: str, system_prompt: str) -> str:
        """Generate documentation in a single API call."""
        return self.openai_client.generate_documentation(content, system_prompt)
    
    def _generate_multiple_chunks(self, files: List[FileInfo], system_prompt: str) -> str:
        """Generate documentation by processing files in chunks."""
        chunks = self.file_processor.chunk_files(files, self.config.max_tokens_per_chunk)
        
        chunk_documentations = []
        for i, chunk in enumerate(chunks):
            console.print(f"[yellow]Processing chunk {i+1}/{len(chunks)}...[/yellow]")
            chunk_content = self.file_processor.format_files_for_documentation(chunk)
            chunk_doc = self.openai_client.generate_documentation(chunk_content, system_prompt)
            chunk_documentations.append(chunk_doc)
        
        # Merge chunks
        console.print("[yellow]Merging chunks...[/yellow]")
        merged_content = "\n\n".join(chunk_documentations)
        
        # Create final merged documentation
        merge_prompt = """You are an expert at merging documentation sections. 
        Create a cohesive, well-structured README.md from the provided documentation chunks.
        Remove duplicates, maintain logical flow, and ensure consistent formatting."""
        
        return self.openai_client.generate_documentation(merged_content, merge_prompt)
    
    def _save_documentation(self, documentation: str, output_path: Path) -> None:
        """Save documentation to file."""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(documentation)
            console.print(f"[green]Documentation saved to: {output_path}[/green]")
        except Exception as e:
            console.print(f"[red]Error saving documentation: {e}[/red]")
            raise
    
    def validate_repository(self, repository_path: Path, additional_skip_patterns: Set[str] = None) -> dict:
        """Validate repository and return statistics."""
        try:
            stats = self.file_processor.get_repository_stats(repository_path, additional_skip_patterns)
            
            # Estimate tokens
            files = stats["files"]
            total_content = self.file_processor.format_files_for_documentation(files)
            estimated_tokens = self.openai_client.count_tokens(total_content)
            estimated_cost = self.openai_client.estimate_cost(estimated_tokens)
            
            return {
                "valid": True,
                "stats": stats,
                "estimated_tokens": estimated_tokens,
                "estimated_cost": estimated_cost,
                "needs_chunking": estimated_tokens > self.config.max_tokens_per_chunk
            }
        except Exception as e:
            return {
                "valid": False,
                "error": str(e)
            } 