"""
Merge all data sources into a single training corpus.
Combines: train.txt, llms-full.txt, and all text files from repos_cloned/
"""
import sys; sys.path.insert(0, '.')
from pathlib import Path

data_dir = Path("data")
out_path = data_dir / "corpus.txt"

extensions = {".py", ".rst", ".md", ".html", ".txt", ".cfg", ".ini",
              ".toml", ".yaml", ".yml", ".json", ".css", ".js", ".bat",
              "Makefile", "dockerignore", "gitignore"}

segments = []

# 1. Existing train.txt
p = data_dir / "train.txt"
if p.exists():
    segments.append(("train.txt", p.read_text(encoding="utf-8")))

# 2. llms-full.txt
p = data_dir / "llms-full.txt"
if p.exists():
    segments.append(("llms-full.txt", p.read_text(encoding="utf-8")))

# 3. All text files from cloned repos
repos = data_dir / "repos_cloned"
if repos.exists():
    for f in sorted(repos.rglob("*")):
        if not f.is_file():
            continue
        if f.suffix in extensions or f.name in extensions:
            try:
                text = f.read_text(encoding="utf-8", errors="replace")
                if len(text) > 50:
                    segments.append((f.relative_to(data_dir).as_posix(), text))
            except Exception:
                pass

print(f"Found {len(segments)} file segments")

# Combine with file markers
lines = []
for name, text in segments:
    lines.append(f"<|file|>{name}")
    lines.append(text)
    lines.append("")

combined = "\n".join(lines)
out_path.write_text(combined, encoding="utf-8")
size_mb = len(combined) / 1e6
print(f"Written: {out_path} ({size_mb:.1f} MB, {len(combined):,} chars)")