import json
import os
import sys
from collections import defaultdict

# ======================
# 配置（可按需调整）
# ======================
TOKEN_DICT_FILE = "token_dict.json"  # 增量token词典（记录token ID和出现次数）
NGRAM_TABLE_FILE = "ngram_table.json"  # 各阶gram表（原始token序列）
MAX_N = 10  # 最大扫描阶数（实际会提前停止）
MIN_FREQ = 1  # 最小计数（至少出现1次才记录）

# ======================
# 工具函数：加载/保存JSON
# ======================
def load_json(path):
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

def save_json(obj, path):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)

# ======================
# 增量Token词典：仅记录token和计数（无抽象ID干扰）
# ======================
def update_token_dict(token, token_dict):
    if token not in token_dict:
        token_dict[token] = {"id": len(token_dict), "count": 1}
    else:
        token_dict[token]["count"] += 1
    return token_dict

# ======================
# 核心：增量扫描N-Gram，返回「是否有新增序列」
# ======================
def scan_ngram(token_stream, n, ngram_table):
    """
    扫描n阶gram，增量更新ngram_table，返回：
    - has_new: 是否发现新的n-gram序列
    - new_grams: 新增的n-gram列表
    - total_grams: 该阶总n-gram列表（含历史）
    """
    n_key = str(n)
    # 初始化当前阶的gram表（如果不存在）
    if n_key not in ngram_table:
        ngram_table[n_key] = defaultdict(int)
    
    # 扫描当前token流的n-gram
    stream_len = len(token_stream)
    if stream_len < n:
        return False, [], dict(ngram_table[n_key])  # 长度不足，无新序列
    
    # 统计当前流的n-gram
    current_grams = defaultdict(int)
    for i in range(stream_len - n + 1):
        gram_tuple = tuple(token_stream[i:i+n])  # 用tuple存原始token序列
        gram_key = " → ".join(gram_tuple)       # 转成"token1→token2→..."格式
        current_grams[gram_key] += 1
    
    # 增量更新到全局表，并判断是否有新增
    has_new = False
    new_grams = []
    total_grams = ngram_table[n_key].copy()
    
    for gram_key, cnt in current_grams.items():
        if cnt < MIN_FREQ:
            continue
        # 新增序列：不在历史表中
        if gram_key not in total_grams:
            has_new = True
            new_grams.append(gram_key)
        # 累加计数
        total_grams[gram_key] += cnt
    
    # 更新全局表
    ngram_table[n_key] = total_grams
    return has_new, new_grams, total_grams

# ======================
# 读取Token文件：按空格分割原始Token（无换行干扰）
# ======================
def read_token_file(file_path):
    if not os.path.exists(file_path):
        print(f"❌ 文件不存在：{file_path}")
        return []
    with open(file_path, "r", encoding="utf-8") as f:
        # 读取所有内容，按空格分割，过滤空字符串
        tokens = [t.strip() for t in f.read().split() if t.strip()]
    print(f"✅ 读取文件 {file_path}：共 {len(tokens)} 个原始Token")
    return tokens

# ======================
# 主逻辑：处理单个Token文件，逐阶扫描直到无新增
# ======================
def process_single_file(file_path, token_dict, ngram_table):
    # 1. 读取原始Token流
    token_stream = read_token_file(file_path)
    if not token_stream:
        return
    
    # 2. 更新Token词典（仅统计，不影响gram扫描）
    for token in token_stream:
        token_dict = update_token_dict(token, token_dict)
    
    # 3. 逐阶扫描2~MAX_N gram，无新增则立即停止
    print(f"\n{'='*60}")
    print(f"开始扫描 {file_path} 的N-Gram（直到无新增序列）")
    print(f"{'='*60}")
    
    last_has_new = True
    for n in range(2, MAX_N + 1):
        if not last_has_new:
            print(f"\n🛑 {n}阶及以上无新增序列，停止扫描")
            break
        
        # 扫描当前阶gram
        has_new, new_grams, total_grams = scan_ngram(token_stream, n, ngram_table)
        last_has_new = has_new
        
        # 直观打印结果
        print(f"\n📊 {n}-Gram 扫描结果：")
        print(f"  - 新增序列数：{len(new_grams)}")
        print(f"  - 累计序列数：{len(total_grams)}")
        
        # 打印新增序列（带计数）
        if new_grams:
            print(f"  - 新增序列（原始Token）：")
            for gram in new_grams:
                print(f"    [{gram}] → 出现 {total_grams[gram]} 次")
        else:
            print(f"  - 无新增序列")

# ======================
# 入口：支持命令行传入多个Token文件
# ======================
if __name__ == "__main__":
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("📝 用法：python3 script.py <token文件1> <token文件2> ...")
        print("示例：python3 script.py ggml.cpp.tokens other.cpp.tokens")
        sys.exit(1)
    
    # 加载历史数据
    token_dict = load_json(TOKEN_DICT_FILE)
    ngram_table = load_json(NGRAM_TABLE_FILE)
    
    # 处理所有传入的文件
    for file_path in sys.argv[1:]:
        process_single_file(file_path, token_dict, ngram_table)
    
    # 保存增量结果
    save_json(token_dict, TOKEN_DICT_FILE)
    save_json(ngram_table, NGRAM_TABLE_FILE)
    
    # 最终汇总
    print(f"\n{'='*60}")
    print("📈 最终汇总（所有文件）：")
    print(f"  - 总Token数：{len(token_dict)}")
    for n in range(2, MAX_N + 1):
        n_key = str(n)
        if n_key not in ngram_table:
            break
        total = len(ngram_table[n_key])
        if total == 0:
            break
        print(f"  - {n}-Gram 累计序列数：{total}")
    print(f"\n💾 结果已保存到：")
    print(f"  - Token词典：{TOKEN_DICT_FILE}")
    print(f"  - N-Gram表：{NGRAM_TABLE_FILE}")