gitllm

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "aiohttp",
#     "rich"
# ]
# ///
import os
import asyncio
import aiohttp
from pathlib import Path
import sys
from fnmatch import fnmatch
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, TextColumn, BarColumn, TimeRemainingColumn
from rich.tree import Tree
import io

# Configuration
DEFAULT_BRANCH = 'main'
GITHUB_TOKEN = os.getenv('GITHUB_TOKEN')
OUTPUT_FILE = 'repo_files.md'
HEADERS = (
	{'Authorization': f'bearer {GITHUB_TOKEN}', 'Content-Type': 'application/json'} if GITHUB_TOKEN else {}
)

# Initialize rich console
console = Console()

# Simplified project file patterns
FILE_PATTERNS = [
	# Key configuration files
	'package.json',
	'svelte.config.js',
	'vite.config.js',
	'vite.config.ts',
	# All relevant file types throughout the project
	'**/*.json',
	'**/*.js',
	'**/*.ts',
	'**/*.svelte',
]


async def fetch_repo_files(repo, branch='main'):
	"""Fetch all matching files from a repository."""
	owner, name = repo.split('/')

	# Check repository and get default branch if needed
	with console.status('[bold green]Checking repository...', spinner='dots'):
		url = f'https://api.github.com/repos/{owner}/{name}'
		async with aiohttp.ClientSession(headers=HEADERS) as session:
			async with session.get(url) as response:
				response.raise_for_status()
				repo_info = await response.json()

		if branch == 'main' and repo_info['default_branch'] != 'main':
			branch = repo_info['default_branch']
			console.print(f'[yellow]Using default branch:[/] [bold cyan]{branch}[/]')

	# Get repository contents using Git Trees API
	with console.status('[bold green]Fetching repository structure...', spinner='dots'):
		tree_url = f'https://api.github.com/repos/{owner}/{name}/git/trees/{branch}?recursive=1'
		async with aiohttp.ClientSession(headers=HEADERS) as session:
			async with session.get(tree_url) as response:
				response.raise_for_status()
				tree_data = await response.json()

		if tree_data.get('truncated', False):
			console.print('[yellow]Warning: Repository tree is truncated. Some files may be missing.[/]')

		# Filter files based on patterns
		matching_files = []
		for item in tree_data.get('tree', []):
			if item['type'] == 'blob' and any(fnmatch(item['path'], pattern) for pattern in FILE_PATTERNS):
				matching_files.append({'path': item['path']})

	console.print(f'Found [bold green]{len(matching_files)}[/] matching files.')
	return matching_files


async def fetch_file_contents(repo, files, branch):
	"""Fetch file contents asynchronously with progress tracking."""
	owner, name = repo.split('/')
	results = []

	with Progress(
		TextColumn('[bold blue]{task.description}'),
		BarColumn(),
		TextColumn('[progress.percentage]{task.percentage:>3.0f}%'),
		TimeRemainingColumn(),
	) as progress:
		overall_task = progress.add_task(f'[cyan]Downloading files...', total=len(files))

		async with aiohttp.ClientSession() as session:
			# Process in batches to avoid rate limits
			batch_size = 10  # Increase batch size slightly for better performance
			for i in range(0, len(files), batch_size):
				batch = files[i : i + batch_size]
				urls = [
					f"https://raw.githubusercontent.com/{owner}/{name}/{branch}/{file['path']}"
					for file in batch
				]
				paths = [file['path'] for file in batch]

				tasks = [fetch_single_file(session, url, path) for url, path in zip(urls, paths)]
				batch_results = await asyncio.gather(*tasks, return_exceptions=True)

				for j, result in enumerate(batch_results):
					file_path = batch[j]['path']
					if isinstance(result, Exception):
						console.print(f'[red]Error fetching {file_path}: {result}[/]')
					else:
						results.append({'path': file_path, 'content': result})

				progress.update(overall_task, advance=len(batch))

	return results


async def fetch_single_file(session, url, path):
	"""Fetch a single file's content."""
	async with session.get(url) as response:
		if response.status != 200:
			raise ValueError(f'Failed to fetch file: {response.status}')
		return await response.text()


def generate_tree_structure(files):
	"""Generate a tree structure for the output markdown using rich."""
	# Build directory tree dictionary
	path_dict = {}
	for file in files:
		parts = file['path'].split('/')
		current = path_dict
		for i, part in enumerate(parts):
			if i == len(parts) - 1:  # File
				if 'files' not in current:
					current['files'] = []
				current['files'].append(part)
			else:  # Directory
				if part not in current:
					current[part] = {}
				current = current[part]

	# Build tree using rich's Tree
	tree = Tree('Repository')

	def build_tree(node, tree_node):
		# Process directories
		for name, contents in sorted([(k, v) for k, v in node.items() if k != 'files']):
			branch = tree_node.add(f'{name}/')
			build_tree(contents, branch)

		# Process files
		for file in sorted(node.get('files', [])):
			tree_node.add(file)

	build_tree(path_dict, tree)

	return tree


async def main():
	# Parse command line arguments
	if len(sys.argv) < 2:
		console.print('[bold red]Error:[/] Please provide a repository name (owner/repo) or URL')
		sys.exit(1)

	# Parse repo from argument (handle URLs too)
	repo_arg = sys.argv[1]
	if repo_arg.startswith('https://'):
		parts = repo_arg.strip('/').split('/')
		if len(parts) >= 4 and parts[2] in ['github.com', 'www.github.com']:
			repo = f'{parts[-2]}/{parts[-1]}'
		else:
			console.print('[bold red]Error:[/] Invalid GitHub URL format')
			sys.exit(1)
	else:
		repo = repo_arg

	# Get output file path (optional second argument)
	output_file = sys.argv[2] if len(sys.argv) > 2 else OUTPUT_FILE

	# Get branch (optional --branch argument)
	branch = DEFAULT_BRANCH
	if '--branch' in sys.argv and sys.argv.index('--branch') + 1 < len(sys.argv):
		branch_idx = sys.argv.index('--branch') + 1
		branch = sys.argv[branch_idx]

	console.print(
		Panel.fit(
			f'[bold cyan]Downloading relevant Svelte files for LLM[/]\n[green]{repo}[/] (branch: [yellow]{branch}[/])',
			border_style='cyan',
		)
	)

	# Use a try-except-finally pattern for better cleanup
	try:
		# Fetch matching files
		files = await fetch_repo_files(repo, branch)
		if not files:
			console.print('[yellow]No matching files found in repository.[/]')
			return

		# Sort files for consistent output
		files.sort(key=lambda f: f['path'])

		# Fetch file contents
		file_contents = await fetch_file_contents(repo, files, branch)
		console.print(f'Successfully fetched [bold green]{len(file_contents)}/{len(files)}[/] files')

		# Generate tree structure
		tree = generate_tree_structure(files)

		# Display the tree in terminal
		console.print(tree)

		# Capture tree as string for markdown
		str_io = io.StringIO()
		console_capture = Console(file=str_io, width=100)
		console_capture.print(tree)
		tree_md = str_io.getvalue()

		# Generate markdown content
		# This handles both owner/repo and URL formats
		repo_name = repo.split('/')[-1]
		md_content = f'# {repo_name} Project Structure\n\n'
		md_content += '## File Tree\n\n'
		md_content += f'```\n{tree_md}\n```\n\n'
		md_content += '## File Contents\n\n'

		# Add file contents
		for file in file_contents:
			file_path = file['path']
			content = file['content']
			ext = file_path.split('.')[-1] if '.' in file_path else ''
			md_content += f'### {file_path}\n\n```{ext}\n{content}\n```\n\n'

		# Write to file
		Path(output_file).write_text(md_content, encoding='utf-8')
		console.print(f'[bold green]Generated:[/] {output_file}')

	except Exception as e:
		console.print(f'[bold red]Error:[/] {e}')
		sys.exit(1)


if __name__ == '__main__':
	asyncio.run(main())