"""
Curate the raw corpus into a small, ultra-high-quality subset.

Splits the corpus into documents, scores each on quality heuristics,
and keeps only the top-tier real source code. Produces data/corpus_curated.txt.

Quality signals:
  - Real source code (NOT RST docs, YAML configs, JSON data, prose, etc.)
  - Balanced delimiters (braces/parens/brackets roughly match)
  - Has function/class definitions at line starts (not embedded in prose)
  - Has comments or docstrings (documentation density)
  - Reasonable length (200–50k chars)
  - Low repetition (unique line ratio)
  - Clean ASCII (no encoding garbage)
  - High code-to-prose ratio
"""

import os
import re

PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CORPUS_PATH = os.path.join(PROJECT_DIR, "data", "corpus.txt")
CURATED_PATH = os.path.join(PROJECT_DIR, "data", "corpus_curated.txt")

# ─── Code structure patterns (must be at line start, not in prose) ───────────
CODE_SIGNATURES = {
    "python":   re.compile(r"^(def |class |import |from \S+ import |if __name__|@|elif |except |finally )", re.M),
    "js_ts":    re.compile(r"^(function |const |let |var |class |export |import |async function|interface |type \w+ =)", re.M),
    "rust":     re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |use |pub struct|pub enum|macro_rules!)", re.M),
    "go":       re.compile(r"^(func |package |import |type \w+ struct|var |const )", re.M),
    "c_cpp":    re.compile(r"^(#include|#define|#ifndef|#ifdef|#if |#endif|typedef |struct \w+|class \w+)", re.M),
    "java":     re.compile(r"^(public |private |protected |class \w+|import |package )", re.M),
}

# ─── Documentation / non-code exclusion patterns ─────────────────────────────
RST_DIRECTIVE = re.compile(r"^\.\. [\w-]+::", re.M)          # .. c:function::
RST_LABEL = re.compile(r"^\.\. _[\w-]+:", re.M)              # .. _faq-cache:
RST_UNDERLINE = re.compile(r"^[-=~^\"']{3,}\s*$", re.M)      # ----, ====, ~~~~
RST_CROSSREF = re.compile(r":[\w]+:`", re.M)                 # :func:`name`
XML_START = re.compile(r"^\s*<\?xml|<!DOCTYPE|<\w+ xmlns", re.M)
MARKDOWN_HEADER = re.compile(r"^#{1,6}\s+\w", re.M)          # # Header

# Prose detection: lines that look like English sentences
PROSE_LINE = re.compile(r"^[A-Z][a-z]+ .* [a-z]+\.$", re.M)  # "The quick brown fox."


def is_documentation(doc: str) -> bool:
    """Check if a document is RST/Markdown documentation rather than code."""
    # Strong RST signals
    if RST_DIRECTIVE.search(doc) or RST_LABEL.search(doc):
        return True
    if RST_CROSSREF.search(doc):
        return True
    # Section underlines (RST/Markdown)
    if RST_UNDERLINE.findall(doc):
        return True
    # XML
    if XML_START.search(doc):
        return True
    # Markdown headers
    if MARKDOWN_HEADER.findall(doc) and len(MARKDOWN_HEADER.findall(doc)) > 3:
        return True
    return False


def detect_language(doc: str) -> str | None:
    """Detect if a document is real source code and which language."""
    if is_documentation(doc):
        return None

    for lang, pat in CODE_SIGNATURES.items():
        matches = pat.findall(doc)
        if len(matches) >= 3:  # at least 3 code-structure keywords
            return lang
    return None


def score_document(doc: str) -> tuple[float, str | None]:
    """Score a document 0.0–1.0 on quality. Returns (score, language)."""
    lines = doc.split("\n")
    n_lines = len(lines)
    if n_lines < 5:
        return 0.0, None

    length = len(doc)
    if length < 200 or length > 100_000:
        return 0.0, None

    # Must be real code
    lang = detect_language(doc)
    if lang is None:
        return 0.0, None

    score = 0.0

    # ─── Base: is real code ──────────────────────────────────────────────────
    score += 0.20

    # ─── Length quality (sweet spot: 500–20000 chars) ────────────────────────
    if 500 <= length <= 20000:
        score += 0.15
    elif 200 <= length <= 50000:
        score += 0.08

    # ─── Delimiter balance ───────────────────────────────────────────────────
    braces = doc.count("{") - doc.count("}")
    parens = doc.count("(") - doc.count(")")
    brackets = doc.count("[") - doc.count("]")
    total_delims = doc.count("{") + doc.count("(") + doc.count("[")
    if total_delims > 0:
        imbalance = abs(braces) + abs(parens) + abs(brackets)
        balance_ratio = 1.0 - (imbalance / max(total_delims, 1))
        score += 0.15 * max(balance_ratio, 0.0)

    # ─── Code structure density ──────────────────────────────────────────────
    struct_count = sum(len(p.findall(doc)) for p in CODE_SIGNATURES.values())
    struct_density = min(struct_count / max(n_lines, 1) * 8, 1.0)
    score += 0.15 * struct_density

    # ─── Comment density (sweet spot: 5–30%) ─────────────────────────────────
    comment_lines = 0
    for line in lines:
        s = line.strip()
        if s.startswith("#") or s.startswith("//") or s.startswith("/*") \
           or s.startswith("*") or s.startswith('"""') or s.startswith("'''") \
           or s.startswith("///") or s.startswith("//!"):
            comment_lines += 1
    comment_ratio = comment_lines / max(n_lines, 1)
    if 0.05 <= comment_ratio <= 0.30:
        score += 0.15
    elif 0.02 <= comment_ratio <= 0.40:
        score += 0.08

    # ─── Low repetition ──────────────────────────────────────────────────────
    unique_lines = len(set(lines))
    unique_ratio = unique_lines / max(n_lines, 1)
    if unique_ratio > 0.6:
        score += 0.10
    elif unique_ratio > 0.4:
        score += 0.05
    else:
        score -= 0.05

    # ─── Clean ASCII ─────────────────────────────────────────────────────────
    non_ascii = sum(1 for c in doc if ord(c) > 127)
    if non_ascii / max(length, 1) < 0.01:
        score += 0.05

    # ─── Indentation quality ─────────────────────────────────────────────────
    indented = sum(1 for l in lines if l.startswith("    ") or l.startswith("\t"))
    if indented > 0 and indented / max(n_lines, 1) > 0.1:
        score += 0.05

    # ─── Penalize high prose ratio ───────────────────────────────────────────
    prose_lines = len(PROSE_LINE.findall(doc))
    prose_ratio = prose_lines / max(n_lines, 1)
    if prose_ratio > 0.15:
        score -= 0.15  # too much prose = documentation

    return min(max(score, 0.0), 1.0), lang


def main():
    print(f"Loading corpus from {CORPUS_PATH}...")
    with open(CORPUS_PATH, "r", encoding="utf-8") as f:
        text = f.read()
    print(f"  Corpus size: {len(text) / 1e6:.1f} MB")

    # Split into documents (3+ consecutive newlines = boundary)
    docs = re.split(r"\n{3,}", text)
    print(f"  Total documents: {len(docs):,}")

    # Score all documents
    print("Scoring documents...")
    scored = []
    for i, doc in enumerate(docs):
        s, lang = score_document(doc)
        if s > 0:
            scored.append((s, i, doc, lang))

    print(f"  Documents passing initial filter: {len(scored):,}")

    # Sort by score descending
    scored.sort(key=lambda x: x[0], reverse=True)

    # Keep top tier: score >= 0.5, target ~3 MB
    # Cap per-language to ensure diversity (no single language > 40% of docs)
    LANG_CAPS = {"c_cpp": 200, "js_ts": 150, "python": 150, "rust": 100, "go": 50, "java": 50}
    lang_counts_so_far = {}
    top_tier = []
    total_size = 0
    target_size = 3_000_000

    for score, idx, doc, lang in scored:
        if score < 0.5:
            break
        cap = LANG_CAPS.get(lang, 100)
        if lang_counts_so_far.get(lang, 0) >= cap:
            continue
        if total_size + len(doc) > target_size:
            break
        top_tier.append((score, idx, doc, lang))
        total_size += len(doc)
        lang_counts_so_far[lang] = lang_counts_so_far.get(lang, 0) + 1

    print(f"\nCurated subset:")
    print(f"  Documents: {len(top_tier):,}")
    print(f"  Total size: {total_size / 1e6:.2f} MB")
    print(f"  Score range: {top_tier[-1][0]:.3f} – {top_tier[0][0]:.3f}")

    # Language distribution
    lang_counts = {}
    for _, _, _, lang in top_tier:
        lang_counts[lang] = lang_counts.get(lang, 0) + 1
    print(f"  Language distribution: {lang_counts}")

    # Write curated corpus
    with open(CURATED_PATH, "w", encoding="utf-8") as f:
        for _, _, doc, _ in top_tier:
            f.write(doc.strip())
            f.write("\n\n\n")

    print(f"\nCurated corpus written to {CURATED_PATH}")
    print(f"  Size: {total_size / 1e6:.2f} MB ({len(top_tier):,} documents)")

    # Show samples
    print("\n" + "=" * 60)
    print("SAMPLE DOCUMENTS (top 5 by score)")
    print("=" * 60)
    for i in range(min(5, len(top_tier))):
        score, idx, doc, lang = top_tier[i]
        print(f"\n--- Score: {score:.3f} | Lang: {lang} | Doc #{idx} | {len(doc)} chars ---")
        print(doc[:600])
        if len(doc) > 600:
            print("...")


if __name__ == "__main__":
    main()