import wikipediaapi
import json
import concurrent.futures
from tqdm import tqdm
import re
import os
import requests
import time

# ==========================================
# 本番用 大規模データ生成スクリプト
# ==========================================
NUM_ARTICLES_TO_FETCH = 5000  # 目標とする高品質データ件数
OUTPUT_FILE = "mindmap_dataset_production.jsonl"
MAX_WORKERS = 10              # スレッド数（多すぎるとAPI制限に引っかかるため10程度）

# ユーザーエージェントを明記（Wikipedia APIのルール）
wiki_wiki = wikipediaapi.Wikipedia(
    user_agent='MindMapStudio_ProductionGen/1.0',
    language='ja',
    extract_format=wikipediaapi.ExtractFormat.WIKI
)

def get_random_titles(count):
    """Wikipedia APIからランダムな記事タイトルを一括取得する"""
    titles = set()
    print(f"ランダムな記事のタイトルを {count} 件収集中...")
    pbar = tqdm(total=count)
    while len(titles) < count:
        # 一度に最大50件まで取得可能
        limit = min(50, count - len(titles))
        url = f"https://ja.wikipedia.org/w/api.php?action=query&list=random&rnnamespace=0&rnlimit={limit}&format=json"
        headers = {'User-Agent': 'MindMapStudio_ProductionGen/1.0 (contact@example.com)'}
        try:
            res = requests.get(url, headers=headers).json()
            for q in res.get("query", {}).get("random", []):
                if q["title"] not in titles:
                    titles.add(q["title"])
                    pbar.update(1)
        except Exception as e:
            print(f"\nAPI Fetch Error: {e}")
            time.sleep(1) # エラー時は少し待機
    pbar.close()
    return list(titles)

def is_valid_article(page):
    """品質チェック（ノイズを除外）"""
    if not page.exists(): return False
    text = page.text
    if len(text) < 1000: return False # 文字数が少なすぎる（スタブ記事）
    
    sections = page.sections
    if len(sections) < 4: return False # 見出しが少なすぎる（構造がない）
    
    return True

def process_section(section, level=1):
    """再帰的にセクションを処理し、ノイズセクションを除去"""
    title = section.title.strip()
    
    # AIの学習に不要なメタセクションを完全除外
    exclude_keywords = ["脚注", "出典", "参考文献", "関連項目", "外部リンク", "注釈", "ギャラリー", "一覧"]
    if any(ex in title for ex in exclude_keywords):
        return "", ""
        
    markdown = f"{'#' * level} {title}\n"
    plain_text = section.text.strip() + "\n" if section.text.strip() else ""
    
    for sub_section in section.sections:
        sub_md, sub_pt = process_section(sub_section, level + 1)
        markdown += sub_md
        plain_text += sub_pt
        
    return markdown, plain_text

def fetch_and_process(title):
    try:
        page = wiki_wiki.page(title)
        if not is_valid_article(page):
            return None
            
        markdown_output = f"# {page.title}\n"
        plain_input = page.summary + "\n"
        
        for section in page.sections:
            sec_md, sec_pt = process_section(section, level=2)
            markdown_output += sec_md
            plain_input += sec_pt
            
        # 改行などのクリーニング
        plain_input = re.sub(r'\n+', '\n', plain_input).strip()
        markdown_output = markdown_output.strip()
        
        # 最終チェック: 構造が浅すぎるものは除外
        if len(plain_input) < 500 or len(markdown_output.split('\n')) < 5:
            return None
            
        return {
            "instruction": "以下の長文から論理構造を抽出し、Markdown形式の目次（マインドマップ）を出力してください。",
            "input": plain_input,
            "output": markdown_output
        }
    except Exception:
        return None

def main():
    print("--- 本番環境用 大規模データ生成パイプライン起動 ---")
    
    # 既存のファイルがあれば削除（やり直し用）
    if os.path.exists(OUTPUT_FILE):
        os.remove(OUTPUT_FILE)
        
    titles = get_random_titles(NUM_ARTICLES_TO_FETCH * 2) # 除外されることを見越して多めに取得
    
    success_count = 0
    print("\n記事のダウンロードと解析（マルチスレッド処理）を開始します...")
    
    with open(OUTPUT_FILE, 'a', encoding='utf-8') as f:
        with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            # 進行状況バーの表示
            results = list(tqdm(executor.map(fetch_and_process, titles), total=len(titles)))
            
            for res in results:
                if res is not None:
                    f.write(json.dumps(res, ensure_ascii=False) + '\n')
                    success_count += 1
                    
                    if success_count >= NUM_ARTICLES_TO_FETCH:
                        break

    print(f"\n完了！ 超高品質なデータセット {success_count} 件を {OUTPUT_FILE} に保存しました。")

if __name__ == "__main__":
    main()