#!/usr/bin/env python
"""
数据集下载脚本 - 下载各领域语料并保存到 database/corpus

用法:
    # 下载单个领域（测试链路）
    python download_data.py news_en --sample-num 10

    # 下载所有领域
    python download_data.py all --sample-num 50

领域列表:
    - news_en:    英文新闻 (CNN/DailyMail)
    - news_zh:    中文新闻 (CLUE news2016zh)
    - academic:   学术论文 (arXiv)
    - code:       代码 (GitHub Python)
    - dialogue:   对话 (OpenAssistant)
    - literature: 文学作品 (Gutenberg)
"""

import os
import re
import json
import random
import argparse
from pathlib import Path
from typing import List, Optional, Dict

# 设置 HuggingFace 镜像（国内加速）
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from datasets import load_dataset


def preprocess_text(text: str, lang: str = "en", min_len: int = 100, max_len: int = 200) -> Optional[str]:
    """
    统一预处理文本：去特殊符号、控制长度

    Args:
        text: 原始文本
        lang: 语言（en/zh）
        min_len: 最小长度（英文=词数，中文=字符数）
        max_len: 最大长度

    Returns:
        预处理后的文本，None 表示不符合长度被过滤
    """
    text = text.strip()
    if not text:
        return None

    if lang == "en":
        # 英文：保留字母、数字、标点、空格
        text = re.sub(r'[^\w\s.,!?;:()"\']', '', text)
        words = text.split()
        if len(words) < min_len or len(words) > max_len:
            return None
        text = ' '.join(words[:max_len]).strip()
    else:
        # 中文：保留汉字、数字、标点
        text = re.sub(r'[^\u4e00-\u9fff0-9.,!?;:()""\']', '', text)
        chars = list(text)
        if len(chars) < min_len or len(chars) > max_len:
            return None
        text = ''.join(chars[:max_len]).strip()

    return text if text else None


def load_news_en(sample_num: int = 50) -> List[str]:
    """加载英文新闻语料 (CNN/DailyMail)"""
    corpus = []
    print("加载 CNN/DailyMail 数据集...")

    try:
        ds = load_dataset("cnn_dailymail", "3.0.0", split="train")
        texts = [item["article"] for item in random.sample(list(ds), min(200, len(ds)))]

        for text in texts:
            processed = preprocess_text(text, lang="en")
            if processed and processed not in corpus:
                corpus.append(processed)
            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 条英文新闻")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_news_zh(sample_num: int = 50) -> List[str]:
    """加载中文新闻语料 (CLUE news2016zh)"""
    corpus = []
    print("加载 CLUE news2016zh 数据集...")

    try:
        ds = load_dataset("clue", "news2016zh", split="train")
        texts = [item["text"] for item in random.sample(list(ds), min(200, len(ds)))]

        for text in texts:
            processed = preprocess_text(text, lang="zh")
            if processed and processed not in corpus:
                corpus.append(processed)
            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 条中文新闻")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_academic(sample_num: int = 50) -> List[str]:
    """加载学术论文语料 (arXiv)"""
    corpus = []
    print("加载 arXiv 数据集...")

    try:
        ds = load_dataset("arxiv_dataset", split="train")
        texts = []

        for item in random.sample(list(ds), min(300, len(ds))):
            categories = item.get("categories", "")
            update_date = item.get("update_date", "")
            year = update_date.split("-")[0] if update_date else "0"

            if categories in ["cs.CL", "cs.LG", "stat.ML"] and year >= "2023":
                texts.append(item.get("abstract", ""))

        for text in texts:
            processed = preprocess_text(text, lang="en")
            if processed and processed not in corpus:
                corpus.append(processed)
            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 篇论文摘要")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_code(sample_num: int = 50) -> List[str]:
    """加载代码语料 (GitHub Python)"""
    corpus = []
    print("加载 codeparrot/github-code 数据集...")

    try:
        ds = load_dataset("codeparrot/github-code", streaming=True, split="train")

        count = 0
        for item in ds:
            if count >= 300:
                break
            count += 1

            lang = item.get("language", "")
            path = item.get("path", "").lower()

            if lang == "Python" and ".py" in path:
                code = item.get("code", "")
                processed = preprocess_text(code, lang="en")
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 段代码")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_dialogue(sample_num: int = 50) -> List[str]:
    """加载对话语料 (OpenAssistant)"""
    corpus = []
    print("加载 OpenAssistant/oasst1 数据集...")

    try:
        ds = load_dataset("OpenAssistant/oasst1", split="train")

        # 收集对话
        conversations = {}
        for item in random.sample(list(ds), min(200, len(ds))):
            msg_id = item.get("message_id", "")
            parent_id = item.get("parent_id", "")
            text = item.get("text", "")
            role = item.get("role", "")

            if parent_id and parent_id in conversations:
                conversations[parent_id].append(f"{role}: {text}")
            else:
                conversations[msg_id] = [f"{role}: {text}"]

        # 合并对话
        for msgs in conversations.values():
            if len(msgs) >= 2:
                text = " ".join(msgs[:4])  # 最多 4 轮
                processed = preprocess_text(text, lang="en")
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 段对话")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_literature(sample_num: int = 50) -> List[str]:
    """加载文学语料 (Gutenberg)"""
    corpus = []
    print("加载 gutenberg_english 数据集...")

    try:
        ds = load_dataset("gutenberg_english", split="train")
        texts = [item["text"] for item in random.sample(list(ds), min(200, len(ds)))]

        for text in texts:
            processed = preprocess_text(text, lang="en")
            if processed and processed not in corpus:
                corpus.append(processed)
            if len(corpus) >= sample_num:
                break

        print(f"  成功提取 {len(corpus)} 段文学文本")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


# 领域加载器映射
DOMAIN_LOADERS = {
    "news_en": load_news_en,
    "news_zh": load_news_zh,
    "academic": load_academic,
    "code": load_code,
    "dialogue": load_dialogue,
    "literature": load_literature,
}


def save_corpus(texts: List[str], domain: str, output_dir: str = "database/corpus") -> None:
    """保存语料到 JSONL 文件"""
    output_path = Path(output_dir) / domain / "texts.jsonl"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w", encoding="utf-8") as f:
        for i, text in enumerate(texts):
            record = {
                "text": text,
                "metadata": {
                    "index": i,
                    "domain": domain,
                }
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    print(f"  已保存至：{output_path}")


def main():
    parser = argparse.ArgumentParser(description="下载语料数据集")
    parser.add_argument(
        "domain",
        type=str,
        choices=list(DOMAIN_LOADERS.keys()) + ["all"],
        help="领域标识，或 'all' 下载所有"
    )
    parser.add_argument(
        "--sample-num",
        type=int,
        default=10,
        help="每个领域采样数量 (默认：10)"
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="database/corpus",
        help="输出目录 (默认：database/corpus)"
    )

    args = parser.parse_args()

    print("=" * 50)
    print("语料数据集下载")
    print("=" * 50)
    print(f"目标领域：{args.domain}")
    print(f"采样数量：{args.sample_num}")
    print(f"输出目录：{args.output_dir}")
    print()

    if args.domain == "all":
        domains = list(DOMAIN_LOADERS.keys())
    else:
        domains = [args.domain]

    for domain in domains:
        print(f"\n[{domain}]")
        loader = DOMAIN_LOADERS[domain]
        texts = loader(args.sample_num)

        if texts:
            save_corpus(texts, domain, args.output_dir)
        else:
            print(f"  警告：未能获取有效语料")

    print("\n" + "=" * 50)
    print("下载完成!")
    print("=" * 50)


if __name__ == "__main__":
    main()