"""
Download ~500M characters of training data from Wikipedia, coding, and cybersecurity sources.
Saves to data/corpus.txt
"""
import os
import sys
import json
import random
import hashlib
import urllib.request
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

DATA_DIR = Path("data")
DATA_DIR.mkdir(exist_ok=True)
OUT_PATH = DATA_DIR / "corpus.txt"
TARGET_CHARS = 500_000_000  # 500M characters

# ── Wikipedia via HuggingFace datasets ──────────────────────────────
def download_wikipedia(target_chars: int) -> str:
    """Download Wikipedia articles via HuggingFace streaming."""
    print("[wiki] Loading Wikipedia from HuggingFace...")
    try:
        from datasets import load_dataset
    except ImportError:
        print("[wiki] 'datasets' not installed, skipping Wikipedia")
        return ""

    ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
    chunks = []
    total = 0

    for article in ds:
        text = article.get("text", "")
        if len(text) < 100:
            continue
        chunks.append(text)
        total += len(text)
        if total >= target_chars:
            break
        if len(chunks) % 10000 == 0:
            print(f"  wiki: {total/1e6:.1f}M chars from {len(chunks)} articles")

    print(f"  wiki: {total/1e6:.1f}M chars from {len(chunks)} articles")
    return "\n\n".join(chunks)


# ── StackOverflow via HuggingFace ───────────────────────────────────
def download_stackoverflow(target_chars: int) -> str:
    """Download StackOverflow Q&A."""
    print("[so] Loading StackOverflow from HuggingFace...")
    try:
        from datasets import load_dataset
    except ImportError:
        return ""

    try:
        ds = load_dataset("lcornell/stackoverflow", split="train", streaming=True)
    except Exception:
        try:
            ds = load_dataset("HuggingFaceFW/stackoverflow", split="train", streaming=True)
        except Exception:
            print("[so] Could not load StackOverflow dataset, skipping")
            return ""

    chunks = []
    total = 0

    for item in ds:
        title = item.get("title", "")
        body = item.get("body", "")
        tags = item.get("tags", "")
        text = f"Title: {title}\nTags: {tags}\n{body}"
        if len(text) < 50:
            continue
        chunks.append(text)
        total += len(text)
        if total >= target_chars:
            break
        if len(chunks) % 10000 == 0:
            print(f"  so: {total/1e6:.1f}M chars from {len(chunks)} items")

    print(f"  so: {total/1e6:.1f}M chars from {len(chunks)} items")
    return "\n\n".join(chunks)


# ── GitHub code via HuggingFace ─────────────────────────────────────
def download_github_code(target_chars: int) -> str:
    """Download GitHub code snippets."""
    print("[code] Loading GitHub code from HuggingFace...")
    try:
        from datasets import load_dataset
    except ImportError:
        return ""

    try:
        ds = load_dataset("codeparrot/github-code", split="train", streaming=True)
    except Exception:
        try:
            ds = load_dataset("HuggingFaceFW/fineweb", name="sample-10BT", split="train", streaming=True)
        except Exception:
            print("[code] Could not load code dataset, skipping")
            return ""

    chunks = []
    total = 0

    for item in ds:
        text = item.get("text", "") or item.get("content", "")
        if len(text) < 50:
            continue
        chunks.append(text)
        total += len(text)
        if total >= target_chars:
            break
        if len(chunks) % 10000 == 0:
            print(f"  code: {total/1e6:.1f}M chars from {len(chunks)} files")

    print(f"  code: {total/1e6:.1f}M chars from {len(chunks)} files")
    return "\n\n".join(chunks)


# ── Cybersecurity text via direct URLs ──────────────────────────────
CYBERSEC_URLS = [
    # NIST publications and guidelines
    "https://nvlpubs.nist.gov/nistpubs/CSWP/NIST.CSWP.29.pdf",
    # Common security resources (plain text mirrors)
    "https://raw.githubusercontent.com/enaqx/awesome-pentest/master/README.md",
    "https://raw.githubusercontent.com/Hack-with-Github/Awesome-Hacking/master/README.md",
    "https://raw.githubusercontent.com/awslabs/aws-security-reference-architecture-examples/main/README.md",
]

def download_cybersecurity(target_chars: int) -> str:
    """Download cybersecurity reference material."""
    print("[cyber] Downloading cybersecurity resources...")
    chunks = []
    total = 0

    for url in CYBERSEC_URLS:
        try:
            print(f"  cyber: fetching {url[:60]}...")
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=30) as resp:
                data = resp.read().decode("utf-8", errors="replace")
            if len(data) > 100:
                chunks.append(data)
                total += len(data)
        except Exception as e:
            print(f"  cyber: failed {url[:50]}: {e}")
        if total >= target_chars:
            break

    # Also try cybersecurity books/guides from Project Gutenberg
    gutenberg_urls = [
        "https://www.gutenberg.org/files/27734/27734-0.txt",  # Cybersecurity-related
    ]
    for url in gutenberg_urls:
        try:
            print(f"  cyber: fetching gutenberg {url[-20:]}...")
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=30) as resp:
                data = resp.read().decode("utf-8", errors="replace")
            if len(data) > 100:
                chunks.append(data)
                total += len(data)
        except Exception as e:
            print(f"  cyber: failed: {e}")
        if total >= target_chars:
            break

    print(f"  cyber: {total/1e6:.1f}M chars from {len(chunks)} sources")
    return "\n\n".join(chunks)


# ── Merge and write ─────────────────────────────────────────────────
def merge_and_write(sources: list[tuple[str, str]], out_path: Path, target: int):
    """Merge sources proportionally and write to file."""
    total_chars = sum(len(s) for _, s in sources if s)
    if total_chars == 0:
        print("ERROR: No data downloaded")
        sys.exit(1)

    print(f"\nTotal available: {total_chars/1e6:.1f}M chars")
    print(f"Target: {target/1e6:.1f}M chars")

    # Build output with file markers
    lines = []
    written = 0
    for name, text in sources:
        if not text:
            continue
        # Take proportionally to fill target
        proportion = len(text) / total_chars
        take_chars = int(target * proportion)
        taken = text[:take_chars]
        lines.append(f"<|file|>{name}")
        lines.append(taken)
        lines.append("")
        written += len(taken)

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


# ── Main ────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print(f"Target: {TARGET_CHARS/1e6:.0f}M characters\n")

    wiki = download_wikipedia(TARGET_CHARS // 2)
    so = download_stackoverflow(TARGET_CHARS // 4)
    code = download_github_code(TARGET_CHARS // 4)
    cyber = download_cybersecurity(TARGET_CHARS // 8)

    sources = [
        ("wikipedia", wiki),
        ("stackoverflow", so),
        ("github_code", code),
        ("cybersecurity", cyber),
    ]

    merge_and_write(sources, OUT_PATH, TARGET_CHARS)
    print("\nDONE — data ready at data/corpus.txt")