"""Download human-written coding text into data/prompts.txt and data/replies.txt.

prompts.txt: what people type to an assistant (OASST user turns, MBPP, CoNaLa).
replies.txt: OASST assistant turns (human-written replies).

Sources (all human-written, permissively licensed; see data/SOURCES.md):
  - OpenAssistant/oasst2  (Apache-2.0)  English prompter and assistant messages from
                                        conversations whose opening prompt is about code
  - google-research-datasets/mbpp (CC-BY-4.0)  crowd-written task descriptions
  - neulab/conala (MIT)  curated programming intents (paired train + test)

Only prose is kept: code blocks, code-looking lines and markdown markup are
dropped, and each message is cut to its first sentences (prompts: 2, replies: 5).
"""

import gzip
import json
import re
from pathlib import Path

import pyarrow.parquet as pq
from huggingface_hub import hf_hub_download

DATA = Path(__file__).resolve().parent.parent / "data"
MAX_WORDS = 24
MAX_PROMPT_SENTENCES = 2
MAX_REPLY_SENTENCES = 5

# Strong signals only; words like "list" or "error" alone are too ambiguous.
CODE_RE = re.compile(
    r"\b(python|javascript|typescript|java|c\+\+|c#|rust|golang|kotlin|swift|php|ruby|perl|haskell|scala|"
    r"sql|mysql|postgres(ql)?|sqlite|mongodb|bash|powershell|shell script|html|css|react|vue|angular|node\.?js|"
    r"django|flask|numpy|pandas|pytorch|tensorflow|regex|regular expression|json|yaml|xml|csv|api|rest api|"
    r"git|github|docker|kubernetes|linux|terminal|command line|compiler|compile|debug|debugging|stack trace|"
    r"exception|segfault|code|coding|program|programming|programmer|script|function|algorithm|recursion|"
    r"recursive|data structure|linked list|binary tree|hash ?map|array|dictionary|variable|loop|class|"
    r"object[- ]oriented|unit test|runtime|syntax|library|framework|database|backend|frontend|webpage|website)\b",
    re.I,
)
CODE_LINE_RE = re.compile(r"[{};=<>]|^\s*(def|class|import|from|for|if|return|#include|const|let|var)\b|\w\(.*\)")
SENT_SPLIT_RE = re.compile(r"(?<=[.?!])\s+")


def clean(text: str, max_sentences: int = MAX_PROMPT_SENTENCES) -> list[str]:
    """Prose sentences from a message, without code or markdown."""
    text = re.sub(r"```.*?(```|$)", " ", text, flags=re.S)
    lines = []
    for ln in text.splitlines():
        ln = re.sub(r"^\s*(#+|[-*+]|\d+[.)])\s+", "", ln)  # headings, bullets, numbered items
        ln = re.sub(r"\*\*|__|\[([^\]]*)\]\([^)]*\)", r"\1", ln)  # bold, links
        if ln.strip() and not CODE_LINE_RE.search(ln):
            lines.append(ln)
    text = " ".join(" ".join(lines).split())
    text = re.sub(r"`([^`]{1,30})`", r"\1", text)  # short inline code -> plain word
    text = re.sub(r"`[^`]*`", " ", text)
    text = re.sub(r"https?://\S+", " ", text)
    if not text.isascii():
        return []
    out = []
    for sent in SENT_SPLIT_RE.split(text)[:max_sentences]:
        words = sent.split()
        if 3 <= len(words) <= MAX_WORDS:
            out.append(" ".join(words))
    return out


def oasst() -> tuple[list[str], list[str]]:
    path = hf_hub_download("OpenAssistant/oasst2", "2023-11-05_oasst2_ready.messages.jsonl.gz", repo_type="dataset")
    rows = [json.loads(line) for line in gzip.open(path, "rt")]
    rows = [
        r for r in rows
        if r["lang"] == "en" and not r["deleted"] and r["review_result"] and not r["synthetic"]
    ]
    coding_trees = {r["message_tree_id"] for r in rows if r["parent_id"] is None and CODE_RE.search(r["text"])}
    prompts, replies = [], []
    for r in rows:
        if r["message_tree_id"] in coding_trees:
            if r["role"] == "prompter":
                prompts += clean(r["text"], MAX_PROMPT_SENTENCES)
            else:
                replies += clean(r["text"], MAX_REPLY_SENTENCES)
    return prompts, replies


def mbpp() -> list[str]:
    out = []
    for split in ["train", "test", "validation", "prompt"]:
        path = hf_hub_download("google-research-datasets/mbpp", f"full/{split}-00000-of-00001.parquet", repo_type="dataset")
        for text in pq.read_table(path, columns=["text"]).column("text").to_pylist():
            out += clean(text)
    return out


def conala() -> list[str]:
    out = []
    for split in ["train", "test"]:
        path = hf_hub_download("neulab/conala", f"data/conala-paired-{split}.json", repo_type="dataset")
        for line in open(path):
            if line.strip():
                out += clean(json.loads(line)["intent"])
    return out


def main() -> None:
    oasst_prompts, oasst_replies = oasst()
    sources = {
        "prompts": [("oasst2_user", oasst_prompts), ("mbpp", mbpp()), ("conala", conala())],
        "replies": [("oasst2_assistant", oasst_replies)],
    }
    seen, report = set(), {}
    DATA.mkdir(exist_ok=True)
    for kind, parts in sources.items():
        lines = []
        for name, got in parts:
            report[name] = len(got)
            for s in got:
                if s.lower() not in seen:
                    seen.add(s.lower())
                    lines.append(s)
        (DATA / f"{kind}.txt").write_text("\n".join(lines) + "\n")
        report[f"{kind}_unique"] = len(lines)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()