#!/usr/bin/env python3
"""
Script to download PyTorch codebase and extract all Python files into a single string.
Uses <bof> and <eof> markers to denote beginning and end of files.
"""

import os
import subprocess
import tempfile
from pathlib import Path
from typing import Generator


def clone_pytorch_repo(temp_dir: str) -> str:
    """Clone the PyTorch repository to a temporary directory."""
    pytorch_dir = os.path.join(temp_dir, "pytorch")
    print("Cloning PyTorch repository...")

    try:
        subprocess.run(
            [
                "git",
                "clone",
                "--depth",
                "1",
                "https://github.com/pytorch/pytorch.git",
                pytorch_dir,
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        print(f"Successfully cloned PyTorch to {pytorch_dir}")
        return pytorch_dir
    except subprocess.CalledProcessError as e:
        print(f"Error cloning repository: {e}")
        print(f"stdout: {e.stdout}")
        print(f"stderr: {e.stderr}")
        raise


def find_python_files(directory: str) -> Generator[Path, None, None]:
    """Recursively find all Python files in the directory."""
    pytorch_path = Path(directory)

    # Skip certain directories that are unlikely to contain useful Python code
    skip_dirs = {
        ".git",
        "__pycache__",
        ".pytest_cache",
        "build",
        "dist",
        "third_party",
        "test",
        "tests",
        "examples",
        "benchmarks",
        "docs",
        "tools/autograd",
        "tools/setup_helpers",
    }

    for py_file in pytorch_path.rglob("*.py"):
        # Check if any parent directory is in skip_dirs
        if any(part in skip_dirs for part in py_file.parts):
            continue

        # Skip if file is too large (> 1MB)
        try:
            if py_file.stat().st_size > 1024 * 1024:
                continue
        except (OSError, IOError):
            continue

        yield py_file


def extract_python_content(pytorch_dir: str) -> str:
    """Extract content from all Python files and format with markers."""
    content_parts = []
    file_count = 0

    print("Extracting Python files...")

    for py_file in find_python_files(pytorch_dir):
        try:
            with open(py_file, "r", encoding="utf-8", errors="ignore") as f:
                file_content = f.read()

            # Get relative path from pytorch directory
            rel_path = py_file.relative_to(Path(pytorch_dir))

            # Format with markers
            formatted_content = (
                f"<bof>{rel_path}<bof>\n{file_content}\n<eof>{rel_path}<eof>\n\n"
            )
            content_parts.append(formatted_content)
            file_count += 1

            if file_count % 100 == 0:
                print(f"Processed {file_count} files...")

        except (IOError, OSError, UnicodeDecodeError) as e:
            print(f"Warning: Could not read {py_file}: {e}")
            continue

    print(f"Successfully processed {file_count} Python files")
    return "".join(content_parts)


def save_to_file(content: str, output_file: str = "pytorch_codebase.txt"):
    """Save the extracted content to a file."""
    print(f"Saving content to {output_file}...")

    with open(output_file, "w", encoding="utf-8") as f:
        f.write(content)

    file_size = os.path.getsize(output_file)
    print(f"Saved {file_size:,} bytes to {output_file}")


def main():
    """Main function to orchestrate the download and extraction process."""
    print("Starting PyTorch codebase extraction...")

    # Create temporary directory
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"Using temporary directory: {temp_dir}")

        try:
            # Clone the repository
            pytorch_dir = clone_pytorch_repo(temp_dir)

            # Extract Python file content
            content = extract_python_content(pytorch_dir)

            # Save to file
            save_to_file(content)

            print("Extraction completed successfully!")

        except Exception as e:
            print(f"Error during extraction: {e}")
            raise


if __name__ == "__main__":
    main()
