#!/usr/bin/env python
"""
补充下载缺失的领域数据
"""

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

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) -> str | 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
        return ' '.join(words[:max_len]).strip()
    return None


def save_corpus(texts: List[str], domain: str) -> bool:
    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_code(sample_num: int = 50) -> List[str]:
    """加载代码语料 - 使用无需认证的数据集"""
    corpus = []
    print("尝试加载 Python 代码...")

    # 尝试 1: cosmopedia openstax (数学/代码相关)
    try:
        ds = load_dataset("HuggingFaceTB/cosmopedia", "openstax", split="train", streaming=True)
        count = 0
        for item in ds:
            if count >= 500:
                break
            count += 1
            text = item.get("text", "")
            # 筛选包含代码示例的文本
            if text and ('import ' in text or 'def ' in text or 'class ' in text or 'print(' in text):
                lines = text.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"  cosmopedia openstax: {len(corpus)} 条")
            return corpus
    except Exception as e:
        print(f"  cosmopedia openstax 失败：{e}")

    # 尝试 2: 使用 algorithmica 数据集
    print("  尝试 algorithmica...")
    try:
        ds = load_dataset("Laurencie/algorithmica", split="train", streaming=True)
        count = 0
        for item in ds:
            if count >= 500:
                break
            count += 1
            code = item.get("code", "") or item.get("content", "")
            if code:
                lines = code.strip().split('\n')
                if 3 <= len(lines) <= 60:
                    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"  algorithmica: {len(corpus)} 条")
            return corpus
    except Exception as e:
        print(f"  algorithmica 失败：{e}")

    # 尝试 3: 使用 leetcode 问题
    print("  尝试 leetcode-problems...")
    try:
        ds = load_dataset("xiaomingl2000/leetcode-medium-questions", split="train", streaming=True)
        count = 0
        for item in ds:
            if count >= 500:
                break
            count += 1
            # 提取问题和解答
            question = item.get("question", "")
            solution = item.get("solution", "")
            if question or solution:
                text = f"{question}\n\n{solution}" if solution else question
                processed = preprocess_text(str(text), lang="en")
                if processed and processed not in corpus:
                    corpus.append(processed)
            if len(corpus) >= sample_num:
                break
        if corpus:
            print(f"  leetcode: {len(corpus)} 条")
            return corpus
    except Exception as e:
        print(f"  leetcode 失败：{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", "physics", "chemistry", "biology", "quantum",
                           "deep learning", "optimization", "data structure", "complexity"]

        count = 0
        for item in ds:
            if count >= 5000:
                break
            count += 1
            text = item.get("text", "")
            title = item.get("title", "")
            # 更严格的筛选
            if any(kw in title.lower() for kw in science_keywords):
                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 main():
    print("=" * 50)
    print("补充下载缺失的领域")
    print("=" * 50)

    # 补充 code (已有 20 条)
    print("\n[code] 补充到 50 条")
    code_path = Path("database/corpus/code/texts.jsonl")
    code_texts = []
    if code_path.exists():
        with open(code_path) as f:
            for line in f:
                code_texts.append(json.loads(line)["text"])
        print(f"  已有 {len(code_texts)} 条")

    if len(code_texts) < 50:
        need = 50 - len(code_texts)
        new_texts = load_code(need)
        for t in new_texts:
            if t not in code_texts:
                code_texts.append(t)
        save_corpus(code_texts[:50], "code")
    else:
        print("  已足够，无需补充")

    # 补充 academic (已有 20 条)
    print("\n[academic] 补充到 50 条")
    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) < 50:
        need = 50 - len(academic_texts)
        new_texts = load_academic(need)
        for t in new_texts:
            if t not in academic_texts:
                academic_texts.append(t)
        save_corpus(academic_texts[:50], "academic")
    else:
        print("  已足够，无需补充")

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


if __name__ == "__main__":
    main()