"""
Document iterators shared by the tokenizer trainer and the shard preparer.

Two source shapes are supported and both scripts accept either.

* A Hub source names a `repo` and is streamed through `datasets`.
* A local source names a `local_root` and walks that file tree. This is what the
  smoke corpus uses, so the pipeline can be exercised end to end without
  downloading 10GB of gated data first.

A local source entry looks like:

    {"local_root": "~/git", "max_bytes": 50000000, "weight": 1.0}
"""

import ast
import json
import os

CODE_EXTENSIONS = (
    ".py", ".js", ".jsx", ".ts", ".tsx", ".go", ".rs", ".java", ".kt",
    ".c", ".h", ".cc", ".cpp", ".hpp", ".m", ".mm", ".swift",
    ".sh", ".zsh", ".bash", ".sql", ".md", ".toml", ".yaml", ".yml",
)

# Vendored, generated, and build output. None of it is worth tokenizing and some
# of it is large enough to eat the whole byte budget on its own.
SKIP_DIRS = {
    "node_modules", "__pycache__", "site-packages", "vendor", "third_party",
    "venv", "dist", "build", "target", "out", "data", "shards", "checkpoints",
    "Pods", "DerivedData", "Carthage",
}

MAX_FILE_BYTES = 1_000_000
HUB_PATH_FIELDS = {
    "bigcode/starcoderdata": "max_stars_repo_path",
}


def quality_ok(text, path="", syntax_text=None) -> bool:
    """
    Cheap, safe quality gate. Parse where a parser is free, otherwise use the
    structural filters that separate written code from generated blobs.

    This is the verifier-first stance applied to pretraining data rather than to
    generations: at 108M parameters, capacity spent modelling minified bundles and
    base64 payloads is capacity not spent on code. Nothing here executes the
    input, which rules out the obvious way a corpus filter becomes a security
    incident.
    """
    if not text or len(text) < 64:
        return False

    lines = text.split("\n")
    longest = max((len(l) for l in lines), default=0)
    if longest > 1000:
        return False                     # minified, bundled, or a data blob on one line
    if len(text) / max(len(lines), 1) > 120:
        return False                     # mean line length past anything hand written

    alnum = sum(c.isalnum() or c.isspace() for c in text[:20000])
    if alnum / min(len(text), 20000) < 0.55:
        return False                     # base64, hex dumps, encoded assets

    syntax_text = text if syntax_text is None else syntax_text
    ext = os.path.splitext(path)[1].lower()
    if ext == ".py":
        try:
            ast.parse(syntax_text)
        except (SyntaxError, ValueError, MemoryError, RecursionError):
            return False
    elif ext == ".json":
        try:
            json.loads(syntax_text)
        except (ValueError, RecursionError):
            return False
    return True


def hub_path_field(src):
    """Resolve the real path column for a known Hub row schema."""
    explicit = src.get("path_field")
    if explicit:
        return explicit
    return HUB_PATH_FIELDS.get(src.get("repo"), "path")


def hub_path_required(src):
    """Return whether a source declared a schema that must carry its path."""
    return bool(
        src.get("path_field")
        or src.get("repo") in HUB_PATH_FIELDS
    )


def hub_syntax_text(src, text):
    """Remove dataset metadata that is not part of the parsed source file."""
    if src.get("repo") != "bigcode/starcoderdata":
        return text
    first, separator, rest = text.partition("\n")
    if separator and first.startswith("<reponame>"):
        return rest
    return text


def iter_local_texts(src):
    """Walk `local_root` in a stable order, yielding decoded text files."""
    root = os.path.expanduser(src["local_root"])
    if not os.path.isdir(root):
        raise FileNotFoundError(f"local_root does not exist: {root}")

    exts = tuple(src.get("extensions", CODE_EXTENSIONS))
    max_bytes = int(src.get("max_bytes", 50_000_000))
    max_file_bytes = int(src.get("max_file_bytes", MAX_FILE_BYTES))
    # The walk order is stable, so skipping the first N bytes yields files the
    # training corpus never saw. That is how the held out eval set is built.
    skip_bytes = int(src.get("skip_bytes", 0))
    gate = bool(src.get("quality_gate", False))

    consumed = 0
    skipped = 0
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = sorted(
            d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")
        )
        for name in sorted(filenames):
            if not name.endswith(exts):
                continue
            path = os.path.join(dirpath, name)
            try:
                size = os.path.getsize(path)
            except OSError:
                continue
            if size == 0 or size > max_file_bytes:
                continue
            if skipped < skip_bytes:
                skipped += size
                continue
            try:
                with open(path, "r", encoding="utf-8") as f:
                    text = f.read()
            except (OSError, UnicodeDecodeError, ValueError):
                continue
            if gate and not quality_ok(text, path):
                continue
            if not text.strip():
                continue
            yield text
            consumed += size
            if consumed >= max_bytes:
                return


def iter_hub_texts(src, text_field_default="content"):
    """Stream a Hugging Face dataset source one text field at a time."""
    from datasets import load_dataset

    ds = load_dataset(
        src["repo"],
        data_dir=src.get("data_dir"),
        name=src.get("name"),
        split=src.get("split", "train"),
        revision=src.get("revision"),
        streaming=True,
    )
    field = src.get("text_field", text_field_default)
    gate = bool(src.get("quality_gate", False))
    path_field = hub_path_field(src)
    for row in ds:
        text = row.get(field)
        if not text:
            continue
        if (
            gate
            and hub_path_required(src)
            and (
                path_field not in row
                or not isinstance(row[path_field], str)
                or not row[path_field]
            )
        ):
            raise ValueError(
                f"{src['repo']}: required Hub path field "
                f"{path_field!r} is missing or empty"
            )
        path = row.get(path_field, "") or ""
        syntax_text = hub_syntax_text(src, text)
        if gate and not quality_ok(text, path, syntax_text=syntax_text):
            continue
        yield text


def source_texts(src, text_field_default="content"):
    """Dispatch a source entry to the right iterator."""
    if src.get("local_root"):
        return iter_local_texts(src)
    return iter_hub_texts(src, text_field_default)


def describe(src) -> str:
    if src.get("local_root"):
        return f"local:{src['local_root']}"
    return f"{src['repo']} {src.get('data_dir') or src.get('name') or ''}".strip()