#!/usr/bin/env python3
"""
repo-map.py — Token-budgeted repository structure map.

Generates a compact, relevance-ranked snapshot of a codebase.
Inspired by Aider's RepoMap but using ripgrep + git instead of tree-sitter.

Usage:
  python3 repo-map.py ~/Unspooled --budget 1024     # 1024 token budget
  python3 repo-map.py ~/Unspooled --budget 2048     # Full detail
  python3 repo-map.py ~/Unspooled --depth 3          # Max 3 dirs deep
"""

import os
import sys
import json
import subprocess
import argparse
from collections import defaultdict
from datetime import datetime, timezone

# Language detection by extension
LANG_MAP = {
    ".kt": "kotlin", ".kts": "kotlin", ".java": "java",
    ".py": "python", ".js": "js", ".ts": "ts", ".tsx": "tsx",
    ".go": "go", ".rs": "rust", ".rb": "ruby", ".php": "php",
    ".swift": "swift", ".c": "c", ".cpp": "cpp", ".h": "c-header",
    ".hpp": "cpp-header", ".cs": "csharp", ".scala": "scala",
    ".sql": "sql", ".sh": "shell", ".bash": "shell", ".zsh": "shell",
    ".yaml": "yaml", ".yml": "yaml", ".json": "json", ".xml": "xml",
    ".md": "markdown", ".rst": "rst", ".toml": "toml",
    ".gradle": "gradle", ".gradle.kts": "gradle",
    ".ktx": "kotlin", ".proto": "protobuf",
    ".css": "css", ".scss": "scss", ".html": "html",
    ".gradle": "gradle-kts",
}

# File type priority for ranking (higher = more relevant for code understanding)
TYPE_PRIORITY = {
    "source": 10,      # .kt, .py, .go, .rs, .java, .ts, .tsx
    "config": 7,       # .yaml, .toml, .json, .gradle
    "build": 8,        # build.gradle, Makefile, Dockerfile
    "schema": 9,       # .proto, .graphql, .sql
    "test": 6,         # *Test*, *Spec*, test_*
    "docs": 4,         # .md, .rst
    "data": 3,         # .csv, .xml
    "other": 2,
}

IGNORE_DIRS = {".git", "node_modules", "build", ".gradle", "target",
               ".idea", ".vscode", "__pycache__", ".gitlab",
               ".github", "venv", ".venv", ".hermes", "bin", "obj",
               ".kotlin", ".cxx", "out", "dist", ".next", ".turbo",
               ".cache", "_site", "public", "coverage", ".terra"}


def get_changed_recently(filepath: str, days: int = 7) -> bool:
    """Check if file was modified in the last N days using git."""
    try:
        r = subprocess.run(
            ["git", "-C", os.path.dirname(filepath), "log", "-1",
             "--format=%ct", "--", filepath],
            capture_output=True, text=True, timeout=5,
        )
        if r.returncode == 0 and r.stdout.strip():
            ts = int(r.stdout.strip())
            age = datetime.now(timezone.utc).timestamp() - ts
            return age < days * 86400
    except (subprocess.TimeoutExpired, FileNotFoundError, ValueError):
        pass
    return False


def classify_file(fname: str, ext: str) -> tuple:
    """Classify a file by type and language."""
    lang = LANG_MAP.get(ext.lower(), "other")
    basename = os.path.basename(fname).lower()

    # Build files
    if basename in ("build.gradle.kts", "build.gradle", "makefile",
                     "dockerfile", "docker-compose.yml", "docker-compose.yaml",
                     "pom.xml", "cmakelists.txt", "cargo.toml", "package.json"):
        return "build", lang
    # Schema files
    if ext in (".proto", ".graphql", ".sql"):
        return "schema", lang
    # Test files
    if "test" in basename or "spec" in basename or basename.startswith("test_"):
        return "test", lang
    # Source files
    if lang in ("kotlin", "java", "python", "go", "rust", "ts", "tsx",
                "js", "c", "cpp", "swift", "scala", "csharp", "ruby", "php"):
        return "source", lang
    # Config files
    if ext in (".yaml", ".yml", ".json", ".toml", ".ini", ".conf", ".properties"):
        return "config", lang
    # Doc files
    if ext in (".md", ".rst"):
        return "docs", lang

    return "other", lang


def build_repo_map(root_path: str, budget_tokens: int = 1024, max_depth: int = 6):
    """Build a token-budgeted repo structure map."""
    root_path = os.path.abspath(root_path)
    if not os.path.isdir(root_path):
        print(f"::error:: not a directory: {root_path}", file=sys.stderr)
        return []

    files = []
    try:
        
        # Find files fast (no wc -l — too slow on large repos)
        result = subprocess.run(
            ["find", root_path, "-type", "f",
             "-not", "-path", "*/.git/*", "-not", "-path", "*/node_modules/*",
             "-not", "-path", "*/build/*", "-not", "-path", "*/target/*",
             "-not", "-path", "*/venv/*", "-not", "-path", "*/__pycache__/*",
             "-not", "-path", "*/.gradle/*", "-not", "-path", "*/dist/*",
             "-not", "-path", "*/public/*", "-not", "-path", "*/.idea/*",
             "-not", "-path", "*/.vscode/*", "-not", "-path", "*/coverage/*",
             "-not", "-path", "*/bin/*", "-not", "-path", "*/obj/*",
             "-not", "-path", "*.next/*", "-not", "-path", "*.turbo/*",
             "-not", "-path", "*/.cache/*",
             "-not", "-path", "*/site-packages/*",
             "-not", "-path", "*/pip/*", "-not", "-path", "*/tests/*",
             "-not", "-name", "*.lock", "-not", "-name", "*.log",
             "-not", "-name", "*.class", "-not", "-name", "*.apk",
             "-not", "-name", "*.png", "-not", "-name", "*.jpg",
             "-not", "-name", "*.jpeg", "-not", "-name", "*.svg",
             "-not", "-name", "*.ico", "-not", "-name", "*.gif",
             "-not", "-name", "*.webp", "-not", "-name", "*.mp4",
             "-not", "-name", "*.so", "-not", "-name", "*.dex",
             "-not", "-name", "package-lock.json"],
            capture_output=True, text=True, timeout=10,
        )
        
        for line in result.stdout.strip().split("\n"):
            if not line.strip():
                continue
            fpath = line.strip()
            ext = os.path.splitext(fpath)[1].lower()
            # Estimate line count by extension (fast, good enough for ranking)
            if ext in ('.py', '.ts', '.tsx', '.js', '.kt', '.java', '.go', '.rs', '.rb', '.cs'):
                line_count = 200
            elif ext in ('.md', '.rst', '.txt'):
                line_count = 80
            elif ext in ('.yaml', '.yml', '.json', '.toml', '.xml'):
                line_count = 50
            elif ext in ('.html', '.css', '.scss'):
                line_count = 120
            elif ext in ('.gradle', '.gradle.kts'):
                line_count = 60
            else:
                line_count = 30
            # Check depth and ignore dirs
            rel = os.path.relpath(fpath, root_path)
            depth = rel.count(os.sep)
            if depth > max_depth:
                continue
            if any(p in rel.split(os.sep) for p in IGNORE_DIRS):
                continue

            ext = os.path.splitext(fpath)[1].lower()
            ftype, lang = classify_file(fpath, ext)
            changed = get_changed_recently(fpath)

            # Score file
            priority = TYPE_PRIORITY.get(ftype, 2)
            score = priority * 10
            if changed:
                score += 15  # Boost recently changed files
            if line_count < 500:
                score += 5   # Slightly prefer medium files over very large ones
            if line_count < 50:
                score -= 3   # Slightly penalize tiny files (boilerplate/LICENCE)
            # Penalize files that are mostly generated
            basename = os.path.basename(fpath).lower()
            if basename in ("license", "licence", "readme.md", ".gitignore"):
                score = 1
            # Skip noise files entirely
            skip_patterns = (
                "package-lock.json", "yarn.lock", "pnpm-lock.yaml",
                ".DS_Store", "*.class", "*.dex"
            )
            if any(basename == p or fpath.endswith(p.replace("*","")) for p in skip_patterns):
                continue
            # Skip database JSON dumps (auto-generated Room DB backups)
            rel_parts = rel.split(os.sep)
            if any(p.startswith("NexStreamDatabase_") or p.startswith("UnspooledDatabase_") for p in rel_parts):
                continue

            files.append({
                "path": rel,
                "lines": line_count,
                "lang": lang,
                "type": ftype,
                "score": score,
                "changed": changed,
            })
    except subprocess.TimeoutExpired:
        print("::warning:: find timed out at 30s", file=sys.stderr)
        return None
    except FileNotFoundError:
        print("::error:: find not available on this system", file=sys.stderr)
        return None

    # Sort by score descending
    files.sort(key=lambda f: f["score"], reverse=True)

    # Select files within token budget
    selected = []
    used_tokens = 0
    for f in files:
        # Estimate: path (~6 tokens) + lines_estimate (~2 tokens) + metadata (~4 tokens)
        entry_tokens = 12
        if used_tokens + entry_tokens > budget_tokens:
            break
        selected.append(f)
        used_tokens += entry_tokens

    return selected


def format_as_tree(files: list, root: str) -> str:
    """Format selected files as a compact directory tree."""
    if not files:
        return "(empty)"

    lines = []
    lines.append(f"# Repo Map — {os.path.basename(root)}")
    lines.append(f"**{len(files)} files, ~{sum(f['lines'] for f in files)} lines**")
    lines.append(f"*Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}*")
    lines.append("")
    lines.append("| Type | Lang | File | Lines |")
    lines.append("|------|------|------|-------|")

    rendered = 0
    for f in files:
        name = f["path"]
        if len(name) > 60:
            name = "..." + name[-57:]
        marker = " 🔄" if f["changed"] else ""
        lines.append(f"| {f['type']:6s} | {f['lang']:8s} | {name}{marker} | {f['lines']:>5d} |")
        rendered += 1

    lines.append("")
    lines.append(f"*🔄 = modified in last 7 days*")
    lines.append(f"*Budget: {rendered} files selected*")

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Token-budgeted repo structure map")
    parser.add_argument("path", nargs="?", default=".", help="Repository root path")
    parser.add_argument("--budget", type=int, default=1024,
                       help="Token budget for map (default: 1024)")
    parser.add_argument("--depth", type=int, default=10,
                       help="Max directory depth (default: 10)")
    parser.add_argument("--json", action="store_true",
                       help="Output raw JSON instead of formatted tree")

    args = parser.parse_args()

    files = build_repo_map(args.path, args.budget, args.depth)
    if files is None:
        sys.exit(1)

    if args.json:
        print(json.dumps(files, indent=2))
    else:
        print(format_as_tree(files, args.path))


if __name__ == "__main__":
    main()