codechunk-cli/codechunk/cli.py

import os
from pathlib import Path
from typing import Optional, List, cast

import click
from rich.console import Console
from rich.panel import Panel
from rich.text import Text

from codechunk.config import Config, load_config
from codechunk.core.chunking import CodeChunker
from codechunk.core.parser import CodeParser
from codechunk.core.formatter import OutputFormatter
from codechunk.core.dependency import DependencyAnalyzer
from codechunk.utils.logger import get_logger

console = Console()
logger = get_logger(__name__)


@click.group()
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose output")
@click.option("--config", "-c", type=click.Path(), help="Path to config file")
@click.pass_context
def main(ctx: click.Context, verbose: bool, config: Optional[str]) -> None:
    ctx.ensure_object(dict)
    ctx.obj["verbose"] = verbose  # type: ignore[index]
    ctx.obj["config_path"] = config  # type: ignore[index]

    if verbose:
        logger.setLevel("DEBUG")


@main.command()
@click.argument("path", type=click.Path(exists=True, file_okay=False, dir_okay=True))
@click.option("--output", "-o", type=click.Path(), help="Output file path")
@click.option("--format", "-f", type=click.Choice(["ollama", "lmstudio", "markdown"]), default="markdown",
              help="Output format")
@click.option("--max-tokens", "-t", type=int, default=8192, help="Maximum token limit")
@click.option("--include", multiple=True, help="File patterns to include")
@click.option("--exclude", multiple=True, help="File patterns to exclude")
@click.pass_context
def generate(ctx: click.Context, path: str, output: Optional[str], format: str,
             max_tokens: int, include: tuple, exclude: tuple) -> None:
    """Generate optimized context bundle for LLM."""
    ctx_obj = cast(dict, ctx.obj)
    config_path = ctx_obj.get("config_path")
    verbose = ctx_obj.get("verbose", False)

    try:
        config = load_config(config_path) if config_path else Config()
        if include:
            config.chunking.include_patterns = list(include)
        if exclude:
            config.chunking.exclude_patterns = list(exclude)

        project_path = Path(path)
        logger.info(f"Analyzing project: {project_path}")

        parser = CodeParser()
        parser.discover_files(project_path, config.chunking.include_patterns,
                              config.chunking.exclude_patterns)

        if verbose:
            console.print(f"[cyan]Found {len(parser.files)} files to process[/cyan]")

        chunks = parser.parse_all()

        chunker = CodeChunker(config.chunking)
        chunks = chunker.chunk_all(chunks)

        analyzer = DependencyAnalyzer()
        analyzer.analyze_dependencies(chunks, parser.files)

        formatter = OutputFormatter(format, max_tokens)
        formatted_output = formatter.format(chunks)

        if output:
            output_path = Path(output)
            output_path.write_text(formatted_output)
            console.print(f"[green]Context written to: {output_path}[/green]")
        else:
            console.print(Panel(formatted_output, title="Generated Context", expand=False))

        console.print(f"\n[cyan]Summary:[/cyan]")
        console.print(f"  - Files processed: {len(parser.files)}")
        console.print(f"  - Chunks generated: {len(chunks)}")
        console.print(f"  - Estimated tokens: {formatter.estimate_tokens(formatted_output)}")

    except Exception as e:
        logger.error(f"Error generating context: {e}")
        console.print(f"[red]Error: {e}[/red]")
        raise click.Abort()


@main.command()
@click.argument("path", type=click.Path(exists=True, file_okay=False, dir_okay=True))
@click.option("--json", is_flag=True, help="Output in JSON format")
@click.pass_context
def analyze(ctx: click.Context, path: str, json: bool) -> None:
    """Analyze codebase and report statistics."""
    ctx_obj = cast(dict, ctx.obj)
    verbose = ctx_obj.get("verbose", False)

    try:
        project_path = Path(path)
        logger.info(f"Analyzing project: {project_path}")

        config = Config()
        parser = CodeParser()
        parser.discover_files(project_path, config.chunking.include_patterns,
                              config.chunking.exclude_patterns)
        chunks = parser.parse_all()

        chunker = CodeChunker(config.chunking)
        chunks = chunker.chunk_all(chunks)

        stats: dict[str, int] = {
            "total_files": len(parser.files),
            "total_chunks": len(chunks),
            "total_lines": sum(c.metadata.line_count for c in chunks),
            "total_functions": sum(1 for c in chunks if c.chunk_type == "function"),
            "total_classes": sum(1 for c in chunks if c.chunk_type == "class"),
        }

        files_by_lang: dict[str, int] = {}
        chunks_by_type: dict[str, int] = {}

        for chunk in chunks:
            lang = chunk.metadata.language
            files_by_lang[lang] = files_by_lang.get(lang, 0) + 1
            chunks_by_type[chunk.chunk_type] = chunks_by_type.get(chunk.chunk_type, 0) + 1

        if json:
            import json as json_module
            full_stats: dict[str, object] = {
                "total_files": stats["total_files"],
                "total_chunks": stats["total_chunks"],
                "total_lines": stats["total_lines"],
                "total_functions": stats["total_functions"],
                "total_classes": stats["total_classes"],
                "files_by_language": files_by_lang,
                "chunks_by_type": chunks_by_type,
            }
            console.print(json_module.dumps(full_stats, indent=2))
        else:
            console.print(Panel(
                Text.from_markup(f"""
[b]Project Analysis[/b]

Total Files: {stats['total_files']}
Total Chunks: {stats['total_chunks']}
Total Lines: {stats['total_lines']}
Total Functions: {stats['total_functions']}
Total Classes: {stats['total_classes']}

[b]Files by Language[/b]
{chr(10).join(f'  - {lang}: {count}' for lang, count in files_by_lang.items())}

[b]Chunks by Type[/b]
{chr(10).join(f'  - {type_}: {count}' for type_, count in chunks_by_type.items())}
                """),
                title="Analysis Results",
                expand=False
            ))

    except Exception as e:
        logger.error(f"Error analyzing project: {e}")
        console.print(f"[red]Error: {e}[/red]")
        raise click.Abort()


@main.command()
@click.pass_context
def version(ctx: click.Context) -> None:
    """Show version information."""
    from codechunk import __version__
    console.print(f"[cyan]CodeChunk CLI v{__version__}[/cyan]")


if __name__ == "__main__":
    main()