#!/usr/bin/env python
"""
补充下载 academic 和 literature 到 200 条 - 宽松模式
"""

import os
import re
import json
from pathlib import Path
from typing import List, Optional

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from datasets import load_dataset


def preprocess_text(text: str, lang: str = "en", min_len: int = 20, max_len: int = 500) -> Optional[str]:
    """预处理文本 - 放宽长度限制"""
    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
        return ' '.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
        return ''.join(chars[:max_len]).strip()


def save_corpus(texts: List[str], domain: str) -> bool:
    """保存语料到 JSONL 文件"""
    if not texts:
        return False
    output_path = Path(f"database/corpus/{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):
            f.write(json.dumps({"text": text, "metadata": {"index": i, "domain": domain}}, ensure_ascii=False) + "\n")
    print(f"  ✅ 已保存：{domain} ({len(texts)} 条)")
    return True


def load_academic(sample_num: int = 200) -> List[str]:
    """加载学术语料 - 宽松模式"""
    corpus = []
    print("加载维基百科文章（宽松模式）...")

    # 更广泛的科学/学术关键词
    science_keywords = [
        "algorithm", "neural", "machine learning", "computer", "math",
        "statistics", "artificial", "physics", "chemistry", "biology", "quantum",
        "deep learning", "optimization", "data", "complexity", "relativity",
        "thermodynamics", "electromagnetism", "genetics", "evolution", "ecosystem",
        "calculus", "topology", "algebra", "geometry", "probability", "enzyme", "protein",
        "dna", "rna", "cell", "molecule", "atom", "electron", "photon", "galaxy", "star",
        "science", "research", "theory", "analysis", "model", "system", "engineering"
    ]

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

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

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

            # 宽松筛选：标题或内容前 300 字符包含关键词即可
            if any(kw in title.lower() or kw in text[:300].lower() for kw in science_keywords):
                processed = preprocess_text(text, lang="en", min_len=20, max_len=500)
                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 = 200) -> List[str]:
    """加载文学语料 - 宽松模式"""
    corpus = []
    print("加载维基百科文章（宽松模式）...")

    # 更广泛的人文/文学关键词
    literature_keywords = [
        "novel", "poetry", "fiction", "literature", "shakespeare",
        "dickens", "austen", "tolkien", "hemingway", "orwell",
        "pride and prejudice", "great gatsby", "1984", "lord of the rings",
        "romeo and juliet", "hamlet", "macbeth", "jane eyre", "wuthering heights",
        "moby dick", "war and peace", "crime and punishment", "ulysses",
        "writer", "author", "book", "story", "character", "plot", "narrative",
        "prose", "verse", "sonnet", "drama", "tragedy", "comedy", "english literature",
        "art", "history", "philosophy", "culture", "music", "painting", "artist"
    ]

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

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

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

            # 宽松筛选
            if any(kw in title.lower() or kw in text[:300].lower() for kw in literature_keywords):
                processed = preprocess_text(text, lang="en", min_len=20, max_len=500)
                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 main():
    print("=" * 60)
    print("补充下载 academic 和 literature 到 200 条")
    print("=" * 60)
    print()

    # 补充 academic
    print("[academic]")
    academic_path = Path("database/corpus/academic/texts.jsonl")
    academic_texts = []
    if academic_path.exists():
        with open(academic_path) as f:
            for line in f:
                academic_texts.append(json.loads(line)["text"])
        print(f"  已有 {len(academic_texts)} 条")

    if len(academic_texts) < 200:
        need = 200 - len(academic_texts)
        print(f"  需要补充 {need} 条")
        new_texts = load_academic(need)
        for t in new_texts:
            if t not in academic_texts:
                academic_texts.append(t)
        save_corpus(academic_texts[:200], "academic")
    else:
        print("  已足够，无需补充")

    print()

    # 补充 literature
    print("[literature]")
    literature_path = Path("database/corpus/literature/texts.jsonl")
    literature_texts = []
    if literature_path.exists():
        with open(literature_path) as f:
            for line in f:
                literature_texts.append(json.loads(line)["text"])
        print(f"  已有 {len(literature_texts)} 条")

    if len(literature_texts) < 200:
        need = 200 - len(literature_texts)
        print(f"  需要补充 {need} 条")
        new_texts = load_literature(need)
        for t in new_texts:
            if t not in literature_texts:
                literature_texts.append(t)
        save_corpus(literature_texts[:200], "literature")
    else:
        print("  已足够，无需补充")

    print()
    print("=" * 60)
    print("完成!")
    print("=" * 60)


if __name__ == "__main__":
    main()