#!/usr/bin/env python3
"""
parse_lean_declarations.py — 从 Lean 源码解析声明库 declarations.jsonl
输入: Lean 源码根目录 (openai_math-main/lean, ~150K .lean 文件)
输出: declarations.jsonl — 每行 {"name": 短名, "kind": ..., "signature": ..., "docstring": ...}

与 README_接入方法.md 第五节口径一致:
  - 正则匹配 (theorem|lemma|def|abbrev|instance|structure|class|inductive|opaque) 名字
  - 签名取到第一个 := 或 where 前
  - docstring 取声明上方紧邻的 /-- ... -/ 块
  - 短名 = 点分名最后一段; 冲突保留首见
"""
import re, json, os, sys, time

LEAN_ROOT = sys.argv[1] if len(sys.argv) > 1 else "/mnt/d/AI_for_Science/openai_math-main/lean"
OUT = sys.argv[2] if len(sys.argv) > 2 else "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/declarations.jsonl"

KINDS = ("theorem", "lemma", "def", "abbrev", "instance", "structure", "class", "inductive", "opaque")
# 声明起始: 行首(允许空白) kind 空格 标识符
DECL_RE = re.compile(
    r'^\s*(theorem|lemma|def|abbrev|instance|structure|class|inductive|opaque)\s+'
    r'([A-Za-z_α-ωΑ-Ω0-9][\wα-ωΑ-Ω.\'’]*)\b'
)
# docstring 块: /-- ... -/ (多行)
DOC_OPEN = re.compile(r'^\s*/--')
DOC_CLOSE = re.compile(r'-/\s*$')

def short_name(full: str) -> str:
    return full.split('.')[-1]

def collect_files(root):
    for dirpath, _, filenames in os.walk(root):
        for fn in filenames:
            if fn.endswith('.lean'):
                yield os.path.join(dirpath, fn)

def parse_file(path):
    """返回该文件内 [(name, kind, signature, docstring), ...]"""
    out = []
    try:
        with open(path, encoding='utf-8', errors='replace') as f:
            lines = f.readlines()
    except Exception:
        return out
    n = len(lines)
    # 预扫描: 每个 docstring 块的结束行 -> 内容
    # 用简单状态机记录最近的 docstring 结束行
    doc_end_line = -1  # 最近一个 docstring 块的结束行索引
    doc_text = ""
    in_doc = False
    doc_buf = []
    doc_end_line = -1

    i = 0
    # 先收集 docstring 结束位置
    doc_ends = {}  # line_idx -> docstring text (不含 /-- 和 -/)
    j = 0
    while j < n:
        line = lines[j]
        if not in_doc and DOC_OPEN.match(line):
            in_doc = True
            doc_buf = []
            # 同行闭合?
            rest = line
            # 去掉开头的 /--
            m = DOC_OPEN.match(line)
            # 收集这行 /-- 之后的内容
            content = line[m.end():]
            if DOC_CLOSE.search(content):
                # 单行 docstring
                content = DOC_CLOSE.sub('', content).strip()
                doc_buf.append(content)
                doc_ends[j] = ' '.join(doc_buf).strip()
                in_doc = False
                doc_buf = []
            else:
                doc_buf.append(content.strip())
        elif in_doc:
            if DOC_CLOSE.search(line):
                content = DOC_CLOSE.sub('', line).strip()
                if content:
                    doc_buf.append(content)
                doc_ends[j] = ' '.join(x for x in doc_buf if x).strip()
                in_doc = False
                doc_buf = []
            else:
                doc_buf.append(line.strip())
        j += 1

    # 扫描声明
    k = 0
    while k < n:
        m = DECL_RE.match(lines[k])
        if not m:
            k += 1
            continue
        kind = m.group(1)
        full_name = m.group(2)
        name = short_name(full_name)
        # docstring: 找声明行上方最近的 docstring 结束行 (空行容忍)
        doc = ""
        for back in range(k-1, max(k-4, -1), -1):
            if back in doc_ends:
                doc = doc_ends[back]
                break
            if lines[back].strip() != "":
                # 遇到非空非 docstring 行，停止
                break
        # 签名: 从声明名之后累积到第一个 := 或 where 或换行
        sig_parts = []
        # 当前行名字之后的内容
        rest = lines[k][m.end():]
        sig_parts.append(rest.rstrip('\n'))
        # 若当前行没有 := / where，继续向下累积
        joined = rest
        kk = k + 1
        if not re.search(r':=|\bwhere\b', joined):
            while kk < n:
                joined += lines[kk]
                sig_parts.append(lines[kk].rstrip('\n'))
                if re.search(r':=|\bwhere\b', lines[kk]):
                    break
                kk += 1
                if kk - k > 40:  # 保护: 签名过长
                    break
        signature = ' '.join(sig_parts)
        signature = re.split(r':=|\bwhere\b', signature)[0].strip()
        # 压缩空白
        signature = re.sub(r'\s+', ' ', signature).strip()
        out.append((name, kind, signature, doc))
        k += 1
    return out

def main():
    t0 = time.time()
    seen = {}  # 短名 -> 首见声明
    files = list(collect_files(LEAN_ROOT))
    total = len(files)
    print(f"[parse] 共 {total} 个 .lean 文件", flush=True)
    count = 0
    for idx, path in enumerate(files):
        for (name, kind, sig, doc) in parse_file(path):
            if name not in seen:
                seen[name] = {"name": name, "kind": kind, "signature": sig, "docstring": doc}
        count += 1
        if count % 20000 == 0:
            print(f"[parse] 进度 {count}/{total} ({time.time()-t0:.0f}s), 声明累计 {len(seen)}", flush=True)
    # 写出
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, 'w', encoding='utf-8') as f:
        for d in seen.values():
            f.write(json.dumps(d, ensure_ascii=False) + '\n')
    print(f"[parse] DONE: {len(seen)} 声明 -> {OUT} ({time.time()-t0:.0f}s)", flush=True)

if __name__ == '__main__':
    main()