#!/usr/bin/env python3
"""
TLDR Code File -- Compressed code summaries for AI agents.

Reads source files and produces compressed summaries that capture
the essential structure (exports, functions, classes, types, imports)
in ~5-10% of the original token count. Agents can read these summaries
instead of full files during review, testing, and planning phases.

Output per file:
  - Imports (grouped by source)
  - Exports list
  - Function/class signatures with parameter types
  - Type definitions
  - Key constants

Usage:
    tldr-codefile src/App.tsx                    # summarize one file
    tldr-codefile src/ --recursive               # summarize all files in directory
    tldr-codefile src/ --recursive --out .tldr/   # write summaries to directory
    tldr-codefile src/ --recursive --stats        # show compression stats
"""

import sys
import os
import re
import json
from pathlib import Path
from collections import defaultdict


# ============================================================
# TypeScript/JavaScript Summarizer
# ============================================================

def summarize_typescript(content, filepath):
    """Extract structural summary from TypeScript/JavaScript."""
    summary = {
        'imports': [],
        'exports': [],
        'functions': [],
        'classes': [],
        'types': [],
        'constants': [],
        'hooks': [],  # React hooks
        'components': [],  # React components
    }

    lines = content.split('\n')

    # --- Imports ---
    import_groups = defaultdict(list)
    for match in re.finditer(
        r'''import\s+(?:(\w+)(?:\s*,\s*)?)?(?:\{([^}]+)\})?\s+from\s+['"]([^'"]+)['"]''',
        content
    ):
        default = match.group(1)
        named = match.group(2)
        source = match.group(3)
        items = []
        if default:
            items.append(default)
        if named:
            items.extend(n.strip().split(' as ')[0].strip() for n in named.split(',') if n.strip())
        import_groups[source].extend(items)

    for source, items in sorted(import_groups.items()):
        summary['imports'].append(f"from '{source}': {', '.join(items)}")

    # --- Exports ---
    # export function/const/class/type/interface
    for match in re.finditer(
        r'export\s+(?:default\s+)?(?:function|const|let|var|class|type|interface|enum)\s+(\w+)',
        content
    ):
        summary['exports'].append(match.group(1))

    # --- Function signatures ---
    for match in re.finditer(
        r'(?:export\s+)?(?:async\s+)?function\s+(\w+)\s*(?:<[^>]+>)?\s*\(([^)]*)\)(?:\s*:\s*([^\s{]+))?',
        content
    ):
        name = match.group(1)
        params = match.group(2).strip()
        ret = match.group(3) or ''
        # Compress params
        if len(params) > 80:
            param_count = params.count(',') + 1
            params = f"...{param_count} params"
        sig = f"{name}({params})"
        if ret:
            sig += f" -> {ret}"
        summary['functions'].append(sig)

    # Arrow function exports
    for match in re.finditer(
        r'export\s+(?:default\s+)?const\s+(\w+)\s*(?::\s*[\w<>[\]|&\s]+)?\s*=\s*(?:async\s+)?\(([^)]*)\)\s*(?::\s*([^\s=>]+))?\s*=>',
        content
    ):
        name = match.group(1)
        params = match.group(2).strip()
        ret = match.group(3) or ''
        if len(params) > 80:
            param_count = params.count(',') + 1
            params = f"...{param_count} params"
        sig = f"{name}({params})"
        if ret:
            sig += f" -> {ret}"
        summary['functions'].append(sig)

    # --- Classes ---
    for match in re.finditer(
        r'(?:export\s+)?class\s+(\w+)(?:\s+extends\s+([\w.]+))?(?:\s+implements\s+([\w,\s]+))?',
        content
    ):
        name = match.group(1)
        extends = match.group(2)
        implements = match.group(3)
        sig = f"class {name}"
        if extends:
            sig += f" extends {extends}"
        if implements:
            sig += f" implements {implements.strip()}"
        summary['classes'].append(sig)

    # --- Type/Interface definitions ---
    for match in re.finditer(
        r'(?:export\s+)?(?:type|interface)\s+(\w+)(?:<[^>]+>)?\s*=?\s*\{?',
        content
    ):
        summary['types'].append(match.group(1))

    # --- Constants ---
    for match in re.finditer(
        r'(?:export\s+)?const\s+(\w+)\s*(?::\s*([\w<>[\]|&\s]+))?\s*=\s*(?!.*=>)',
        content
    ):
        name = match.group(1)
        type_ann = match.group(2)
        # Skip if it's a function (already captured above)
        if name not in [f.split('(')[0] for f in summary['functions']]:
            sig = name
            if type_ann:
                sig += f": {type_ann.strip()}"
            summary['constants'].append(sig)

    # --- React hooks (use* pattern) ---
    for match in re.finditer(
        r'(?:export\s+)?(?:function|const)\s+(use\w+)\s*(?:<[^>]+>)?\s*\(([^)]*)\)',
        content
    ):
        name = match.group(1)
        params = match.group(2).strip()
        if len(params) > 60:
            params = f"...{params.count(',') + 1} params"
        summary['hooks'].append(f"{name}({params})")

    # --- React components (PascalCase function returning JSX) ---
    for match in re.finditer(
        r'(?:export\s+)?(?:default\s+)?(?:function|const)\s+([A-Z]\w+)\s*(?:<[^>]+>)?\s*\(([^)]*)\)',
        content
    ):
        name = match.group(1)
        props = match.group(2).strip()
        if name not in summary['exports']:
            continue
        if len(props) > 60:
            props = "props"
        summary['components'].append(f"<{name} {props}/>")

    return summary


# ============================================================
# Python Summarizer
# ============================================================

def summarize_python(content, filepath):
    """Extract structural summary from Python."""
    summary = {
        'imports': [],
        'exports': [],
        'functions': [],
        'classes': [],
        'types': [],
        'constants': [],
        'hooks': [],
        'components': [],
    }

    # --- Imports ---
    for match in re.finditer(r'^from\s+([\w.]+)\s+import\s+(.+)', content, re.MULTILINE):
        module = match.group(1)
        names = [n.strip().split(' as ')[0].strip() for n in match.group(2).split(',')]
        summary['imports'].append(f"from {module}: {', '.join(names)}")

    for match in re.finditer(r'^import\s+([\w.]+)', content, re.MULTILINE):
        summary['imports'].append(f"import {match.group(1)}")

    # --- Functions ---
    for match in re.finditer(
        r'^(?:async\s+)?def\s+(\w+)\s*\(([^)]*)\)(?:\s*->\s*(\S+))?',
        content, re.MULTILINE
    ):
        name = match.group(1)
        params = match.group(2).strip()
        ret = match.group(3) or ''
        if len(params) > 80:
            param_count = params.count(',') + 1
            params = f"...{param_count} params"
        sig = f"def {name}({params})"
        if ret:
            sig += f" -> {ret}"
        if not name.startswith('_'):
            summary['exports'].append(name)
        summary['functions'].append(sig)

    # --- Classes ---
    for match in re.finditer(
        r'^class\s+(\w+)(?:\(([^)]*)\))?',
        content, re.MULTILINE
    ):
        name = match.group(1)
        bases = match.group(2) or ''
        sig = f"class {name}"
        if bases:
            sig += f"({bases.strip()})"
        summary['classes'].append(sig)
        if not name.startswith('_'):
            summary['exports'].append(name)

    # --- Constants (ALL_CAPS) ---
    for match in re.finditer(r'^([A-Z][A-Z_0-9]+)\s*=\s*(.+)', content, re.MULTILINE):
        name = match.group(1)
        value = match.group(2).strip()[:50]
        summary['constants'].append(f"{name} = {value}")

    return summary


# ============================================================
# Output Formatting
# ============================================================

def format_summary(summary, filepath, original_lines):
    """Format a file summary as compact markdown."""
    parts = []
    parts.append(f"## `{filepath}` ({original_lines} lines)")
    parts.append("")

    if summary['imports']:
        parts.append(f"**Imports:** {len(summary['imports'])}")
        for imp in summary['imports'][:10]:
            parts.append(f"  - {imp}")
        if len(summary['imports']) > 10:
            parts.append(f"  - ... +{len(summary['imports'])-10} more")
        parts.append("")

    if summary['exports']:
        parts.append(f"**Exports:** {', '.join(summary['exports'])}")
        parts.append("")

    if summary['components']:
        parts.append("**Components:**")
        for comp in summary['components']:
            parts.append(f"  - {comp}")
        parts.append("")

    if summary['hooks']:
        parts.append("**Hooks:**")
        for hook in summary['hooks']:
            parts.append(f"  - {hook}")
        parts.append("")

    if summary['classes']:
        parts.append("**Classes:**")
        for cls in summary['classes']:
            parts.append(f"  - {cls}")
        parts.append("")

    if summary['functions']:
        parts.append("**Functions:**")
        for func in summary['functions']:
            parts.append(f"  - {func}")
        parts.append("")

    if summary['types']:
        parts.append(f"**Types:** {', '.join(summary['types'])}")
        parts.append("")

    if summary['constants']:
        parts.append("**Constants:**")
        for const in summary['constants'][:10]:
            parts.append(f"  - {const}")
        if len(summary['constants']) > 10:
            parts.append(f"  - ... +{len(summary['constants'])-10} more")
        parts.append("")

    return '\n'.join(parts)


# ============================================================
# CLI
# ============================================================

def main():
    args = sys.argv[1:]

    if not args or '--help' in args:
        print(__doc__)
        sys.exit(0)

    recursive = '--recursive' in args
    show_stats = '--stats' in args
    out_dir = None

    # Parse --out
    if '--out' in args:
        idx = args.index('--out')
        if idx + 1 < len(args):
            out_dir = Path(args[idx + 1])
            out_dir.mkdir(parents=True, exist_ok=True)

    # Get target path
    target = None
    for a in args:
        if not a.startswith('--') and a not in (out_dir and str(out_dir) or ''):
            target = Path(a)
            break

    if not target:
        print("Usage: tldr-codefile <file-or-directory>")
        sys.exit(1)

    # Collect files
    files = []
    if target.is_file():
        files = [target]
    elif target.is_dir() and recursive:
        skip = {'node_modules', '.git', 'dist', 'build', '__pycache__', '.lab', '.session'}
        for f in sorted(target.rglob('*')):
            if any(s in f.parts for s in skip):
                continue
            if f.suffix in ('.ts', '.tsx', '.js', '.jsx', '.py', '.rs') and f.is_file():
                files.append(f)
    elif target.is_dir():
        for f in sorted(target.iterdir()):
            if f.suffix in ('.ts', '.tsx', '.js', '.jsx', '.py', '.rs') and f.is_file():
                files.append(f)

    if not files:
        print(f"No source files found in {target}")
        sys.exit(1)

    # Summarize
    total_original = 0
    total_summary = 0
    all_summaries = []

    for filepath in files:
        try:
            content = filepath.read_text(encoding='utf-8', errors='ignore')
        except Exception:
            continue

        line_count = content.count('\n') + 1
        total_original += len(content)

        if filepath.suffix in ('.ts', '.tsx', '.js', '.jsx', '.mjs'):
            summary = summarize_typescript(content, filepath)
        elif filepath.suffix == '.py':
            summary = summarize_python(content, filepath)
        else:
            continue

        formatted = format_summary(summary, filepath, line_count)
        total_summary += len(formatted)
        all_summaries.append(formatted)

        if out_dir:
            out_file = out_dir / (filepath.stem + '.tldr.md')
            out_file.write_text(formatted)

    # Output
    output = '\n---\n\n'.join(all_summaries)

    if out_dir:
        index_path = out_dir / 'INDEX.md'
        index_path.write_text(f"# TLDR Code Summaries\n\n{output}")
        print(f"Wrote {len(files)} summaries to {out_dir}/", file=sys.stderr)

    if show_stats:
        ratio = (total_summary / total_original * 100) if total_original > 0 else 0
        print(f"\nCompression Stats:", file=sys.stderr)
        print(f"  Files:        {len(files)}", file=sys.stderr)
        print(f"  Original:     {total_original:,} chars", file=sys.stderr)
        print(f"  Summary:      {total_summary:,} chars", file=sys.stderr)
        print(f"  Compression:  {ratio:.1f}% of original", file=sys.stderr)
    elif not out_dir:
        print(output)


if __name__ == "__main__":
    main()