#!/usr/bin/env python
"""
下载所有领域的数据集（每个领域 50 条样本）

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

用法:
    python database/download_all_domains.py
"""

import os
import re
import json
from pathlib import Path
from typing import List, Optional, Callable, 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 = 30, max_len: int = 300) -> 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 save_corpus(texts: List[str], domain: str, output_dir: str = "database/corpus") -> bool:
    """保存语料到 JSONL 文件"""
    if not texts:
        print(f"  ⚠️ 警告：没有语料可保存")
        return False

    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} ({len(texts)} 条)")
    return True


# ==================== 各领域加载函数 ====================

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", trust_remote_code=True)

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

            text = item["article"]
            processed = preprocess_text(text, lang="en", min_len=30, max_len=300)
            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 备用：csl)"""
    corpus = []
    print("加载中文新闻数据集...")

    # 尝试 1: CLUE csl (中文科学文献)
    try:
        ds = load_dataset("clue", "csl", split="train", streaming=True)

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

            text = item.get("text", "") or item.get("abstract", "")
            if text:
                processed = preprocess_text(text, lang="zh", min_len=50, max_len=200)
                if processed and processed not in corpus:
                    corpus.append(processed)
            if len(corpus) >= sample_num:
                break

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

    # 尝试 2: 使用多语言维基百科
    print("  使用备用方案：维基百科中文文章...")
    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.zh", split="train", streaming=True)

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

            text = item.get("text", "")
            if text:
                processed = preprocess_text(text, lang="zh", min_len=50, max_len=200)
                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]:
    """
    加载学术论文语料

    使用维基百科科学类文章作为可靠数据源
    """
    corpus = []
    print("加载学术语料（维基百科科学文章）...")

    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)

        # 科学相关的关键词
        science_keywords = [
            "algorithm", "neural network", "machine learning",
            "computer science", "mathematics", "statistics",
            "artificial intelligence", "data science", "physics",
            "chemistry", "biology", "quantum", "relativity"
        ]

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

            text = item.get("text", "")
            title = item.get("title", "")

            # 筛选科学相关文章
            if any(kw in title.lower() or kw in text[:500].lower() for kw in science_keywords):
                processed = preprocess_text(text, lang="en", min_len=30, max_len=300)
                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]:
    """加载代码语料 (KosmosCode 或其他 Python 数据集)"""
    corpus = []
    print("加载 Python 代码数据集...")

    # 尝试 1: KosmosCode
    try:
        from datasets import load_dataset
        ds = load_dataset("HuggingFaceTB/cosmopedia", "python", split="train", streaming=True)

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

            code = item.get("text", "") or item.get("code", "")
            if code:
                lines = code.strip().split('\n')
                if 5 <= len(lines) <= 80:
                    code_clean = '\n'.join(line[:120] for line in lines[:40])
                    if code_clean and code_clean not in corpus:
                        corpus.append(code_clean)

            if len(corpus) >= sample_num:
                break

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

    # 尝试 2: 使用 bigcode 数据集
    print("  使用备用方案：bigcode/the-stack...")
    try:
        ds = load_dataset("bigcode/the-stack", data_dir="data/python", split="train", streaming=True)

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

            code = item.get("content", "")
            if code:
                lines = code.strip().split('\n')
                if 5 <= len(lines) <= 80:
                    code_clean = '\n'.join(line[:120] for line in lines[:40])
                    if code_clean and code_clean not in corpus:
                        corpus.append(code_clean)

            if len(corpus) >= sample_num:
                break

        if corpus:
            print(f"  成功提取 {len(corpus)} 段 Python 代码（the-stack）")
            return corpus
    except Exception as e:
        print(f"  the-stack 加载失败：{e}")

    # 尝试 3: 使用简单的 Python 代码示例
    print("  使用备用方案：code alpaca...")
    try:
        ds = load_dataset("QingyiSi/Code-Alpaca-Code-Instruction-Following", split="train", streaming=True)

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

            code = item.get("code", "") or item.get("output", "")
            if code:
                lines = code.strip().split('\n')
                if 3 <= len(lines) <= 50:
                    code_clean = '\n'.join(line[:120] for line in lines[:30])
                    if code_clean and code_clean not in corpus:
                        corpus.append(code_clean)

            if len(corpus) >= sample_num:
                break

        if corpus:
            print(f"  成功提取 {len(corpus)} 段 Python 代码（code-alpaca）")
            return corpus
    except Exception as e:
        print(f"  code-alpaca 加载失败：{e}")

    return corpus


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

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

        # 构建对话树
        messages = {}
        for item in ds:
            msg_id = item.get("message_id", "")
            parent_id = item.get("parent_id", "")
            text = item.get("text", "")
            role = item.get("role", "")

            if parent_id not in messages:
                messages[parent_id] = []
            messages[parent_id].append({
                "id": msg_id,
                "role": role,
                "text": text
            })

        # 提取对话
        for root_msgs in messages.values():
            if len(root_msgs) >= 2:
                # 合并对话
                dialogue_parts = []
                for msg in root_msgs[:4]:  # 最多 4 轮
                    dialogue_parts.append(f"{msg['role']}: {msg['text']}")

                dialogue = " ".join(dialogue_parts)
                processed = preprocess_text(dialogue, lang="en", min_len=30, max_len=300)
                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]:
    """加载文学语料 (English Wikipedia 作为备用)"""
    corpus = []
    print("加载文学语料...")

    # 尝试 1: 使用 pile-cc 或其他文学类数据集
    try:
        ds = load_dataset("Z-Code/BookCorpus", split="train", streaming=True)

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

            text = item.get("text", "")
            if text:
                processed = preprocess_text(text, lang="en", min_len=30, max_len=300)
                if processed and processed not in corpus:
                    corpus.append(processed)

            if len(corpus) >= sample_num:
                break

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

    # 尝试 2: 使用英文维基百科中的人文/文学类文章
    print("  使用备用方案：维基百科英文文章（人文类）...")
    try:
        ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)

        # 人文/文学相关的关键词
        literature_keywords = [
            "novel", "poetry", "fiction", "literature", "shakespeare",
            "dickens", "austen", "tolkien", "hemingway", "orwell",
            "pride and prejudice", "great gatsby", "1984", "lord of the rings"
        ]

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

            text = item.get("text", "")
            title = item.get("title", "")

            # 筛选文学相关文章
            if any(kw in title.lower() or kw in text[:500].lower() for kw in literature_keywords):
                processed = preprocess_text(text, lang="en", min_len=30, max_len=300)
                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: Dict[str, Callable] = {
    "news_en": load_news_en,
    "news_zh": load_news_zh,
    "academic": load_academic,
    "code": load_code,
    "dialogue": load_dialogue,
    "literature": load_literature,
}


def main():
    print("=" * 60)
    print("数据集下载 - 完整语料库（每个领域 50 条样本）")
    print("=" * 60)
    print()

    domains = list(DOMAIN_LOADERS.keys())
    sample_num = 50

    print(f"目标领域：{domains}")
    print(f"每个领域采样数：{sample_num}")
    print(f"输出目录：database/corpus/")
    print()

    results = {}

    for domain in domains:
        print(f"\n{'='*40}")
        print(f"[{domain}]")
        print(f"{'='*40}")

        loader = DOMAIN_LOADERS[domain]
        texts = loader(sample_num)

        if texts:
            success = save_corpus(texts, domain)
            results[domain] = len(texts) if success else 0
        else:
            print(f"  ⚠️ 警告：未能获取有效语料")
            results[domain] = 0

    # 汇总
    print("\n" + "=" * 60)
    print("下载完成! 汇总:")
    print("=" * 60)

    for domain, count in results.items():
        status = "✅" if count > 0 else "❌"
        print(f"  {status} {domain}: {count} 条")

    total = sum(results.values())
    print(f"\n总计：{total} 条")

    print()
    print("下一步：提取模型表示")
    print("  python model/extract_representations.py")


if __name__ == "__main__":
    main()