src/chunk_corpus.py

"""
Chunk the curated corpus into searchable code units.

Splits code into semantically meaningful chunks:
  - Python: functions, classes, top-level blocks
  - JS/TS: functions, classes, export blocks
  - Rust: fn, impl, struct, enum, trait blocks
  - Go: func, type, struct blocks
  - C/C++: function definitions, struct/typedef blocks

Each chunk gets:
  - id: unique identifier
  - language: detected language
  - name: extracted name (function/class name)
  - type: function/class/struct/etc
  - code: the raw code text
  - filepath: synthetic path (derived from doc index)
  - start_line, end_line: line range within the doc

Output: data/chunks.jsonl
"""

import json
import os
import re

PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CURATED_PATH = os.path.join(PROJECT_DIR, "data", "corpus_curated.txt")
CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl")

# ─── Language detection ──────────────────────────────────────────────────────
LANG_PATTERNS = {
    "python": re.compile(r"^(def |class |import |from \S+ import |if __name__|@)", re.M),
    "js_ts":  re.compile(r"^(function |const |let |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 |type \w+ struct|import )", re.M),
    "c_cpp":  re.compile(r"^(#include|#define|#ifndef|#if |#endif|typedef |struct \w+|class \w+|void |int |char |size_t |ngx_)", re.M),
}

# ─── Chunk patterns per language ─────────────────────────────────────────────
# Each pattern matches the START of a chunk. We split on these.
CHUNK_STARTS = {
    "python": re.compile(r"^(def |class |@|if __name__)", re.M),
    "js_ts":  re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =)", re.M),
    "rust":   re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |macro_rules!|pub struct|pub enum)", re.M),
    "go":     re.compile(r"^(func |type \w+ struct|type \w+ interface)", re.M),
    "c_cpp":  re.compile(r"^(static |void |int |char |size_t |ngx_|typedef |struct \w+ \{|#define|#if|#ifndef)", re.M),
}

# ─── Name extraction patterns ────────────────────────────────────────────────
NAME_PATTERNS = [
    (re.compile(r"def (\w+)"), "function"),
    (re.compile(r"class (\w+)"), "class"),
    (re.compile(r"fn (\w+)"), "function"),
    (re.compile(r"pub fn (\w+)"), "function"),
    (re.compile(r"struct (\w+)"), "struct"),
    (re.compile(r"enum (\w+)"), "enum"),
    (re.compile(r"trait (\w+)"), "trait"),
    (re.compile(r"impl (\w+)"), "impl"),
    (re.compile(r"func (\w+)"), "function"),
    (re.compile(r"type (\w+) struct"), "struct"),
    (re.compile(r"function (\w+)"), "function"),
    (re.compile(r"typedef struct (\w+)"), "typedef"),
    (re.compile(r"#define (\w+)"), "macro"),
    (re.compile(r"(ngx_\w+)\s*\("), "function"),
]


def detect_language(doc: str) -> str | None:
    for lang, pat in LANG_PATTERNS.items():
        if len(pat.findall(doc)) >= 2:
            return lang
    return None


def extract_name(code: str) -> tuple[str, str]:
    """Extract the name and type from a code chunk."""
    for pat, typ in NAME_PATTERNS:
        m = pat.search(code)
        if m:
            return m.group(1), typ
    return "unknown", "block"


def chunk_document(doc: str, lang: str, doc_idx: int) -> list[dict]:
    """Split a document into chunks based on language-specific patterns."""
    lines = doc.split("\n")
    n_lines = len(lines)

    # Find all chunk start positions
    start_pat = CHUNK_STARTS.get(lang)
    if start_pat is None:
        # Fallback: treat whole doc as one chunk
        return [{
            "id": f"doc_{doc_idx}_chunk_0",
            "language": lang,
            "name": "block",
            "type": "block",
            "code": doc,
            "filepath": f"src/doc_{doc_idx}.txt",
            "start_line": 1,
            "end_line": n_lines,
        }]

    starts = [(m.start(), m.group()) for m in start_pat.finditer(doc)]

    if not starts:
        # No pattern matches — treat whole doc as one chunk
        return [{
            "id": f"doc_{doc_idx}_chunk_0",
            "language": lang,
            "name": "block",
            "type": "block",
            "code": doc,
            "filepath": f"src/doc_{doc_idx}.txt",
            "start_line": 1,
            "end_line": n_lines,
        }]

    # Add doc start if first match isn't at position 0
    if starts[0][0] > 0:
        starts.insert(0, (0, ""))

    chunks = []
    for i, (start_pos, _) in enumerate(starts):
        end_pos = starts[i + 1][0] if i + 1 < len(starts) else len(doc)
        chunk_code = doc[start_pos:end_pos].strip()

        # Skip tiny chunks (< 30 chars)
        if len(chunk_code) < 30:
            continue

        # Skip huge chunks (> 8000 chars — split them)
        if len(chunk_code) > 8000:
            # Split on blank lines
            sub_parts = re.split(r"\n\n+", chunk_code)
            for j, part in enumerate(sub_parts):
                if len(part.strip()) < 30:
                    continue
                name, typ = extract_name(part)
                start_line = doc[:start_pos].count("\n") + 1 + sum(p.count("\n") + 2 for p in sub_parts[:j])
                chunks.append({
                    "id": f"doc_{doc_idx}_chunk_{i}_{j}",
                    "language": lang,
                    "name": name,
                    "type": typ,
                    "code": part.strip(),
                    "filepath": f"src/doc_{doc_idx}.txt",
                    "start_line": start_line,
                    "end_line": start_line + part.count("\n"),
                })
            continue

        name, typ = extract_name(chunk_code)
        start_line = doc[:start_pos].count("\n") + 1

        chunks.append({
            "id": f"doc_{doc_idx}_chunk_{i}",
            "language": lang,
            "name": name,
            "type": typ,
            "code": chunk_code,
            "filepath": f"src/doc_{doc_idx}.txt",
            "start_line": start_line,
            "end_line": start_line + chunk_code.count("\n"),
        })

    return chunks


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

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

    # Chunk each document
    print("Chunking documents...")
    all_chunks = []
    lang_counts = {}
    type_counts = {}

    for i, doc in enumerate(docs):
        doc = doc.strip()
        if len(doc) < 50:
            continue
        lang = detect_language(doc)
        if lang is None:
            continue

        chunks = chunk_document(doc, lang, i)
        for chunk in chunks:
            all_chunks.append(chunk)
            lang_counts[lang] = lang_counts.get(lang, 0) + 1
            type_counts[chunk["type"]] = type_counts.get(chunk["type"], 0) + 1

    print(f"\nTotal chunks: {len(all_chunks):,}")
    print(f"  By language: {lang_counts}")
    print(f"  By type: {type_counts}")

    # Size distribution
    sizes = [len(c["code"]) for c in all_chunks]
    sizes.sort()
    print(f"  Chunk size: min={sizes[0]}, median={sizes[len(sizes)//2]}, max={sizes[-1]}, mean={sum(sizes)//len(sizes)}")

    # Write chunks
    with open(CHUNKS_PATH, "w", encoding="utf-8") as f:
        for chunk in all_chunks:
            f.write(json.dumps(chunk) + "\n")

    print(f"\nChunks written to {CHUNKS_PATH}")

    # Show a few samples
    print("\n" + "=" * 60)
    print("SAMPLE CHUNKS")
    print("=" * 60)
    for i in [0, len(all_chunks) // 4, len(all_chunks) // 2, len(all_chunks) - 1]:
        c = all_chunks[i]
        print(f"\n--- {c['id']} | {c['language']} | {c['type']} | {c['name']} ---")
        print(f"    File: {c['filepath']}:{c['start_line']}-{c['end_line']}")
        print(f"    Size: {len(c['code'])} chars")
        print(c["code"][:300])
        if len(c["code"]) > 300:
            print("...")


if __name__ == "__main__":
    main()