import os
import sys
import argparse
import re

def search(query, path='.', include=None, exclude=None, is_regex=False):
    """
    Search for a string or regex in files.
    """
    results = []
    
    # Compile regex if needed
    if is_regex:
        try:
            pattern = re.compile(query)
        except re.error as e:
            print(f"Error: Invalid regex: {e}")
            return
    
    for root, dirs, files in os.walk(path):
        # Skip some common directories
        if any(d in root for d in ['.git', 'node_modules', 'dist', '__pycache__', '.gemini']):
            continue
            
        for file in files:
            # Filter by extension
            if include and not any(file.endswith(ext) for ext in include):
                continue
            if exclude and any(file.endswith(ext) for ext in exclude):
                continue
                
            file_path = os.path.join(root, file)
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    for i, line in enumerate(f, 1):
                        found = False
                        if is_regex:
                            if pattern.search(line):
                                found = True
                        else:
                            if query in line:
                                found = True
                                
                        if found:
                            results.append({
                                'file': file_path,
                                'line': i,
                                'content': line.strip()
                            })
            except Exception as e:
                # print(f"Error reading {file_path}: {e}")
                continue
                
    return results

def main():
    # Force UTF-8 output for Windows
    import io
    if sys.stdout.encoding != 'utf-8':
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

    parser = argparse.ArgumentParser(description="Clean and reliable code search tool.")
    parser.add_argument("query", help="String or regex to search for")
    parser.add_argument("--path", default=".", help="Directory to search in")
    parser.add_argument("--include", nargs="+", help="Only include files with these extensions (e.g. .vue .ts)")
    parser.add_argument("--exclude", nargs="+", help="Exclude files with these extensions")
    parser.add_argument("--regex", action="store_true", help="Treat query as a regular expression")
    
    args = parser.parse_args()
    
    # Standardize extensions to start with dot
    include = [ext if ext.startswith('.') else f'.{ext}' for ext in args.include] if args.include else None
    exclude = [ext if ext.startswith('.') else f'.{ext}' for ext in args.exclude] if args.exclude else None
    
    results = search(args.query, args.path, include, exclude, args.regex)
    
    if not results:
        print("No results found.")
        return
        
    # Group results by file
    from collections import defaultdict
    grouped = defaultdict(list)
    for res in results:
        grouped[res['file']].append(res)
        
    for file_path, matches in grouped.items():
        print(f"\n\033[1;34m{file_path}\033[0m")
        for m in matches:
            print(f"  \033[1;32m{m['line']}:\033[0m {m['content']}")

if __name__ == "__main__":
    main()