import os
import json
import time
import requests
import wikipediaapi
from tqdm import tqdm

# LM Studioのローカルサーバー設定
LM_STUDIO_URL = "http://127.0.0.1:1234/v1/chat/completions"

# 教師モデルへのプロンプト（システム命令）
SYSTEM_PROMPT = """あなたは超優秀な情報整理AI（マインドマップ職人）です。
ユーザーから提供された長文テキストを深く読み解き、内容を論理的に構造化して、マインドマップ用のMarkdownを出力してください。

【厳格なルール】
1. 出力は必ずMarkdownの階層構造（#、##、###、および箇条書き - ）のみとすること。
2. 一番上の階層（#）は、そのテキストの主題（メインテーマ）を1つだけ記述すること。
3. 下位の階層（##、###）で、大項目・中項目を整理すること。
4. 具体的な要素や詳細な情報は、箇条書き（- ）を用いて末端の枝として表現すること。
5. Markdown以外の説明文、挨拶、感想、「以下に示します」などの前置きは一切出力してはならない（厳禁）。
6. 「概要」や「利用状況」など、テキストにない見出しを適当にでっち上げないこと。テキストに沿った見出しを作ること。
"""

USER_INSTRUCTION = "以下の長文から論理構造を抽出し、Markdown形式の目次（マインドマップ）を出力してください。"

def generate_teacher_output(text):
    """LM StudioのAPIを叩いて、教師モデルにマインドマップを作らせる"""
    payload = {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"{USER_INSTRUCTION}\n\n【対象テキスト】\n{text}"}
        ],
        "temperature": 0.2, # 創造性よりも正確性を重視
        "max_tokens": -1
    }
    
    try:
        response = requests.post(LM_STUDIO_URL, json=payload, timeout=300)
        if response.status_code == 200:
            result = response.json()
            return result['choices'][0]['message']['content'].strip()
        else:
            print(f"API Error: {response.status_code}")
            return None
    except Exception as e:
        print(f"Connection Error: {e}")
        return None

def fetch_random_wikipedia_text(count=10):
    """Wikipediaからランダムな記事のテキストを取得する（見出し構造は使わない）"""
    wiki_wiki = wikipediaapi.Wikipedia(
        user_agent='MindMapStudio_DistillationGen/1.0',
        language='ja',
        extract_format=wikipediaapi.ExtractFormat.WIKI
    )
    
    print(f"Wikipediaからランダムに {count} 件の記事タイトルを取得中...")
    titles = set()
    while len(titles) < count:
        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_DistillationGen/1.0'}
        try:
            res = requests.get(url, headers=headers).json()
            for q in res.get("query", {}).get("random", []):
                titles.add(q["title"])
        except Exception:
            time.sleep(1)
            
    print("記事のテキストをダウンロード中...")
    texts = []
    for title in tqdm(list(titles)):
        try:
            page = wiki_wiki.page(title)
            if page.exists() and len(page.text) > 500:
                # 最初の2000文字程度を学習用テキストとする（長すぎるとAPIが遅くなるため）
                text = page.text[:2000]
                texts.append(text)
        except Exception:
            pass
            
    return texts

def generate_distilled_dataset(num_samples=10, output_file="mindmap_dataset_distilled.jsonl"):
    print("=== 教師モデルによるデータセット蒸留（Knowledge Distillation）を開始 ===")
    
    texts = fetch_random_wikipedia_text(num_samples)
    print(f"合計 {len(texts)} 件の有効なテキストを取得しました。")
    print("LM Studio (Teacher) に推論させています...")
    
    with open(output_file, 'w', encoding='utf-8') as f:
        for text in tqdm(texts):
            teacher_md = generate_teacher_output(text)
            if teacher_md:
                # AIがMarkdownのコードブロック(```markdown)を含めてしまった場合は除去
                teacher_md = teacher_md.replace("```markdown\n", "").replace("```\n", "").replace("```", "").strip()
                
                data = {
                    "instruction": USER_INSTRUCTION,
                    "input": text,
                    "output": teacher_md
                }
                f.write(json.dumps(data, ensure_ascii=False) + '\n')
                f.flush()
            else:
                print("\n警告: APIは成功しましたが、出力(content)が空でした。")
                
    print(f"完了！蒸留データセットを {output_file} に保存しました。")

if __name__ == "__main__":
    # 本番用データセットの生成（500件）
    generate_distilled_dataset(500)