#!/usr/bin/env python3
# description: Search file contents using ripgrep
# Args: <pattern> [--path=<dir>] [--mode=files_with_matches|content|count] [--glob=<pattern>] [--type=<type>] [-i] [-n] [-A <n>] [-B <n>] [-C <n>] [--multiline] [--head=<n>] [--offset=<n>]

import sys
import os
import subprocess
import argparse

parser = argparse.ArgumentParser(prog='file_grep', add_help=False)
parser.add_argument('pattern')
parser.add_argument('--path', default='')
parser.add_argument('--mode', default='files_with_matches',
                    choices=['files_with_matches', 'content', 'count'])
parser.add_argument('--glob', default='')
parser.add_argument('--type', default='')
parser.add_argument('-i', action='store_true')
parser.add_argument('-n', action='store_true', default=True)
parser.add_argument('-A', type=int, default=None)
parser.add_argument('-B', type=int, default=None)
parser.add_argument('-C', type=int, default=None)
parser.add_argument('--multiline', action='store_true')
parser.add_argument('--head', type=int, default=0)
parser.add_argument('--offset', type=int, default=0)

args = parser.parse_args()

rg = 'rg'
search_path = args.path or os.getcwd()

cmd = ['rg', '--no-config', '--color=never']

if args.mode == 'files_with_matches':
    cmd.append('-l')
elif args.mode == 'count':
    cmd.append('--count')
elif args.mode == 'content':
    if args.n:
        cmd.append('-n')

if args.i:
    cmd.append('-i')
if args.multiline:
    cmd += ['-U', '--multiline-dotall']
if args.glob:
    cmd += ['-g', args.glob]
if args.type:
    cmd += ['-t', args.type]
if args.C is not None:
    cmd += ['-C', str(args.C)]
if args.A is not None:
    cmd += ['-A', str(args.A)]
if args.B is not None:
    cmd += ['-B', str(args.B)]

cmd += ['-e', args.pattern, search_path]

result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode == 1:
    print("no matches")
    sys.exit(0)
if result.returncode > 1:
    print(f"error: {result.stderr.strip() or f'rg exited with {result.returncode}'}")
    sys.exit(1)

lines = result.stdout.splitlines()

if args.offset:
    lines = lines[args.offset:]
if args.head:
    lines = lines[:args.head]

if args.mode == 'content':
    formatted = []
    for line in lines:
        if line == '--':
            formatted.append(line)
            continue
        # file:lineno:content or file-lineno-content (context lines)
        for sep in (':', '-'):
            parts = line.split(sep, 2)
            if len(parts) == 3 and parts[1].isdigit():
                formatted.append(f"{parts[0]}:{parts[1]:>5}| {parts[2]}")
                break
        else:
            formatted.append(line)
    lines = formatted

print('\n'.join(lines))
