import wikipediaapi
import json
import time
from tqdm import tqdm

def extract_structure(sections, level=1):
    """
    再帰的にWikipediaのセクション（目次とテキスト）を抽出し、
    正解データ（Markdown）と入力データ（ベタ打ち長文）のペアを作成する。
    """
    markdown_output = ""
    flat_input = ""
    
    for s in sections:
        # 見出しを除外する設定などもあるが、基本的に見出し自体は不要（役に立たないもの）もあるため除外も検討できる
        # 参考文献や関連項目などは除外
        if s.title in ["参考文献", "関連項目", "外部リンク", "脚注", "出典"]:
            continue
            
        # Markdownの構造データ（Output用）
        # H1, H2, H3 などの見出しレベルを付与
        markdown_output += f"{'#' * level} {s.title}\n"
        
        # 概要を箇条書きで追加（各セクションの最初の1文などを要約として入れることも可能）
        # ここではシンプルに見出しのみ、または少しのテキストを入れる
        if s.text:
            first_sentence = s.text.split("。")[0] + "。"
            if len(first_sentence) > 5:
                markdown_output += f"- {first_sentence}\n"
            
        # ベタ打ち長文（Input用）
        # 見出し（s.title）を隠して、テキストだけを連結する
        if s.text:
            flat_input += s.text + "\n"
            
        # サブセクションを再帰的に処理
        child_md, child_flat = extract_structure(s.sections, level + 1)
        markdown_output += child_md
        flat_input += child_flat
        
    return markdown_output, flat_input

def generate_dataset(num_pages=10, output_file="mindmap_dataset.jsonl"):
    # Wikipedia APIの設定 (日本語)
    wiki = wikipediaapi.Wikipedia('MindMapStudio/1.0', 'ja')
    
    print(f"Generating dataset from {num_pages} random Wikipedia articles...")
    
    # 完全にランダムにページを取得するためのハック
    # Wikipediaの「おまかせ表示」APIを利用してランダムなタイトルを取得できるが、
    # ここではテスト用に特定のカテゴリやキーワードからページを取得するか、有名な記事を指定する。
    sample_topics = [
        "量子力学", "人工知能", "日本", "宇宙", "徳川家康", 
        "相対性理論", "コンピュータ", "インターネット", "哲学", "経済学",
        "ブラックホール", "細胞", "DNA", "地球", "太陽系"
    ]
    
    dataset = []
    
    for topic in tqdm(sample_topics[:num_pages]):
        page = wiki.page(topic)
        if not page.exists():
            continue
            
        # ページの最初の要約部分（サマリー）
        summary_md = f"# {page.title}\n- {page.summary.split('。')[0]}。\n"
        summary_flat = page.summary + "\n"
        
        # セクション構造の抽出
        body_md, body_flat = extract_structure(page.sections, level=2)
        
        final_markdown = summary_md + body_md
        final_flat = summary_flat + body_flat
        
        # テキストが短すぎる場合はスキップ
        if len(final_flat) < 100:
            continue
            
        # JSONLフォーマットに整形
        data_row = {
            "instruction": "以下の長文から論理構造を抽出し、Markdown形式の目次（マインドマップ）を出力してください。",
            "input": final_flat.strip(),
            "output": final_markdown.strip()
        }
        dataset.append(data_row)
        time.sleep(1) # APIへの負荷軽減
        
    # JSONLファイルとして保存
    with open(output_file, "w", encoding="utf-8") as f:
        for row in dataset:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
            
    print(f"\nSuccessfully generated {len(dataset)} examples!")
    print(f"Saved to {output_file}")

if __name__ == "__main__":
    generate_dataset(num_pages=5)