"""
Download a multi-language code corpus from GitHub for base pretraining.

Downloads tarballs of several popular repositories across Python, JS/TS,
Rust, Go, C/C++, and Java. Extracts source files and concatenates them
into a single training corpus.
"""

import io
import os
import tarfile
import time
import requests

DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
RAW_DIR = os.path.join(DATA_DIR, "raw")
CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt")

# Repos to download (user/repo, branch, max_bytes_per_repo)
REPOS = [
    ("python/cpython", "main", 8_000_000),       # Python, C
    ("microsoft/TypeScript", "main", 8_000_000),  # TypeScript
    ("rust-lang/rust", "master", 8_000_000),      # Rust
    ("golang/go", "master", 8_000_000),           # Go
    ("nginx/nginx", "master", 5_000_000),         # C
    ("json-iterator/go", "master", 3_000_000),    # Go
    ("pytorch/pytorch", "main", 8_000_000),       # C++, Python, CUDA
    ("tokio-rs/tokio", "master", 5_000_000),      # Rust
    ("numpy/numpy", "main", 5_000_000),           # Python, C
    ("axios/axios", "v1.x", 3_000_000),           # JavaScript
    ("lodash/lodash", "master", 3_000_000),       # JavaScript
    ("redis/redis", "unstable", 5_000_000),       # C
]

# File extensions to include
CODE_EXTENSIONS = {
    ".py", ".js", ".ts", ".mjs", ".jsx", ".tsx",
    ".rs", ".go", ".c", ".h", ".cpp", ".cc", ".cxx",
    ".hpp", ".hxx", ".cu", ".cuh", ".java",
    ".sh", ".bash", ".yml", ".yaml", ".toml",
    ".cfg", ".ini", ".json", ".xml", ".sql",
    ".md", ".txt", ".rst",
}

# Directories to skip
SKIP_DIRS = {
    "test", "tests", "testing", "__pycache__", ".git",
    "node_modules", "vendor", "third_party", "thirdparty",
    "dist", "build", "target", ".github", "docs",
    "benchmark", "benchmarks", "examples", "example",
    "fixtures", "testdata", "test_data", "mocks",
}


def download_repo_tarball(user_repo: str, branch: str) -> bytes | None:
    """Download a repository tarball from GitHub."""
    url = f"https://codeload.github.com/{user_repo}/tar.gz/refs/heads/{branch}"
    print(f"  Downloading {user_repo}@{branch}...")
    try:
        r = requests.get(url, timeout=120, stream=True)
        if r.status_code != 200:
            # Try tags endpoint
            url2 = f"https://codeload.github.com/{user_repo}/tar.gz/refs/tags/{branch}"
            r = requests.get(url2, timeout=120, stream=True)
            if r.status_code != 200:
                print(f"    FAILED (status {r.status_code})")
                return None
        content = r.content
        print(f"    Downloaded {len(content) / 1e6:.1f} MB")
        return content
    except Exception as e:
        print(f"    ERROR: {e}")
        return None


def extract_code_files(tarball_bytes: bytes, max_bytes: int) -> list[str]:
    """Extract code files from a tarball, up to max_bytes total."""
    files = []
    total = 0
    try:
        with tarfile.open(fileobj=io.BytesIO(tarball_bytes), mode="r:gz") as tar:
            for member in tar.getmembers():
                if not member.isfile():
                    continue
                path = member.name
                parts = path.split("/")
                # Skip unwanted directories
                if any(skip in parts for skip in SKIP_DIRS):
                    continue
                # Check extension
                _, ext = os.path.splitext(path)
                if ext.lower() not in CODE_EXTENSIONS:
                    continue
                # Skip very large files
                if member.size > 500_000:
                    continue
                try:
                    f = tar.extractfile(member)
                    if f is None:
                        continue
                    content = f.read()
                    try:
                        text = content.decode("utf-8", errors="ignore")
                    except Exception:
                        continue
                    files.append(text)
                    total += len(text)
                    if total >= max_bytes:
                        break
                except Exception:
                    continue
    except Exception as e:
        print(f"    Extraction error: {e}")
    print(f"    Extracted {len(files)} files, {total / 1e6:.1f} MB of code")
    return files


def main():
    os.makedirs(RAW_DIR, exist_ok=True)
    os.makedirs(DATA_DIR, exist_ok=True)

    all_code = []

    for user_repo, branch, max_bytes in REPOS:
        tarball = download_repo_tarball(user_repo, branch)
        if tarball is None:
            continue
        files = extract_code_files(tarball, max_bytes)
        all_code.extend(files)
        time.sleep(1)  # Be polite

    # Write corpus
    print(f"\nTotal files: {len(all_code)}")
    total_size = sum(len(f) for f in all_code)
    print(f"Total corpus size: {total_size / 1e6:.1f} MB")

    with open(CORPUS_PATH, "w", encoding="utf-8") as f:
        for i, code in enumerate(all_code):
            # Add document separator
            f.write(code)
            f.write("\n<|endoftext|>\n")

    print(f"Corpus written to {CORPUS_PATH}")
    return total_size


if __name__ == "__main__":
    main()