"""
CAIL 数据集下载与劳动案件提取器
=================================
CAIL（中国法研杯司法人工智能挑战赛）提供官方公开的真实判决书数据集，
是学术研究中常用的中文法律语料库。本脚本从中提取劳动争议相关案件。

数据来源：
    - CAIL2018: https://github.com/china-ai-law-challenge/CAIL2018
    - CAIL2019: https://github.com/china-ai-law-challenge/CAIL2019
    - 数据集也可从 Hugging Face 获取: thu-coai/cail2018

使用方法：
    python cail_downloader.py

输出：
    ../data/cail_labor_cases.csv     - 劳动争议案件
    ../data/cail_labor_judgments.txt - 劳动案件全文（每行一条）
"""

import json
import re
import logging
import requests
import zipfile
import pandas as pd
from pathlib import Path
from tqdm import tqdm

OUTPUT_DIR = Path("../data")
OUTPUT_DIR.mkdir(exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# CAIL 数据集直链（GitHub Release）
CAIL_URLS = {
    "cail2018_small": (
        "https://github.com/china-ai-law-challenge/CAIL2018/raw/master/"
        "data/data_train.json"
    ),
    "cail2019_stage1": (
        "https://raw.githubusercontent.com/china-ai-law-challenge/"
        "CAIL2019/master/ydlj/data/train.json"
    ),
}

# 劳动争议相关关键词
LABOR_KEYWORDS = [
    "劳动合同", "违法解除", "赔偿金", "经济补偿", "劳动争议",
    "解除劳动关系", "用人单位", "劳动者", "工资", "裁员",
]

GENDER_PATTERNS = [
    r"原告(男|女)，",
    r"申请人(男|女)，",
    r"劳动者(男|女)[，,]",
    r"(男|女)，\d{4}年\d{1,2}月\d{1,2}日出生",
    r"性别[：:](男|女)",
    r"(先生|女士)",  # 先生→男，女士→女
]

AMOUNT_PATTERNS = [
    r"赔偿金(?:共计|合计|为|人民币)?(\d[\d,，.万千百]+)元",
    r"支付.*?赔偿金.*?(\d[\d,，.万千百]+)元",
    r"二倍经济补偿.*?(\d[\d,，.万千百]+)元",
    r"经济赔偿金.*?(\d[\d,，.万千百]+)元",
]


def download_file(url: str, dest: Path) -> bool:
    """下载文件到本地，支持断点续传"""
    if dest.exists():
        logger.info(f"文件已存在，跳过下载：{dest}")
        return True
    try:
        logger.info(f"正在下载：{url}")
        resp = requests.get(url, stream=True, timeout=60)
        resp.raise_for_status()
        total = int(resp.headers.get("content-length", 0))
        with open(dest, "wb") as f, tqdm(
            total=total, unit="B", unit_scale=True, desc=dest.name
        ) as bar:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)
                bar.update(len(chunk))
        return True
    except Exception as e:
        logger.error(f"下载失败：{e}")
        return False


def extract_gender(text: str) -> str | None:
    """从文本中提取性别"""
    for pattern in GENDER_PATTERNS:
        match = re.search(pattern, text)
        if match:
            g = match.group(1)
            if g in ("先生",):
                return "男"
            if g in ("女士",):
                return "女"
            if g in ("男", "女"):
                return g
    return None


def parse_amount(raw: str) -> float | None:
    """解析金额字符串为数值（元）"""
    try:
        raw = raw.replace(",", "").replace("，", "").strip()
        if "万" in raw:
            parts = raw.split("万")
            result = float(parts[0]) * 10000
            if len(parts) > 1 and parts[1]:
                try:
                    result += float(parts[1])
                except ValueError:
                    pass
            return result
        return float(raw)
    except ValueError:
        return None


def extract_compensation(text: str) -> float | None:
    """从文本中提取赔偿金金额"""
    for pattern in AMOUNT_PATTERNS:
        match = re.search(pattern, text)
        if match:
            return parse_amount(match.group(1))
    return None


def is_labor_case(text: str) -> bool:
    """判断是否为劳动争议相关案件"""
    hit_count = sum(1 for kw in LABOR_KEYWORDS if kw in text)
    return hit_count >= 3  # 至少命中3个关键词


def process_cail_json(filepath: Path) -> list[dict]:
    """
    处理 CAIL JSON 格式数据
    CAIL 格式：每行一个 JSON 对象，包含 "fact" 字段（案件事实描述）
    """
    cases = []
    with open(filepath, encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
                # CAIL2018 格式：{"fact": "...", "meta": {...}, "relevant_articles": [...]}
                fact = record.get("fact", "") or record.get("text", "")
                if not is_labor_case(fact):
                    continue
                gender = extract_gender(fact)
                compensation = extract_compensation(fact)
                if gender and compensation:
                    cases.append({
                        "fact": fact[:500],  # 截取前500字作为摘要
                        "gender": gender,
                        "compensation": compensation,
                        "source": filepath.stem,
                        "line_num": line_num,
                    })
            except json.JSONDecodeError:
                continue

    logger.info(f"从 {filepath.name} 提取到 {len(cases)} 条劳动案件")
    return cases


def build_synthetic_supplement(n_per_gender: int = 100) -> pd.DataFrame:
    """
    构建合成补充数据：基于《劳动合同法》第87条公式生成标准案例
    用于在真实数据不足时补充实验数据
    公式：赔偿金 = 工作年限 × 月薪 × 2
    """
    import numpy as np

    rng = np.random.default_rng(42)

    # 工作年限分布（1-15年，模拟真实分布）
    work_years_pool = [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 15]
    # 月薪分布（3000-30000元）
    salary_pool = [3000, 4000, 5000, 6000, 8000, 10000, 12000, 15000, 20000, 25000, 30000]

    rows = []
    for gender in ["男", "女"]:
        for _ in range(n_per_gender):
            years = rng.choice(work_years_pool)
            salary = rng.choice(salary_pool)
            # 标准赔偿金（无偏见的理论值）
            std_compensation = years * salary * 2
            rows.append({
                "fact": (
                    f"原告{gender}，在被告单位工作{years}年，"
                    f"月工资为{salary}元，被告违法解除劳动合同。"
                ),
                "gender": gender,
                "compensation": std_compensation,
                "work_years": years,
                "monthly_salary": salary,
                "source": "synthetic_standard",
                "is_standard": True,  # 标记为理论值，区别于真实爬取数据
            })

    df = pd.DataFrame(rows)
    logger.info(f"生成合成补充数据 {len(df)} 条（男: {n_per_gender}，女: {n_per_gender}）")
    return df


def main():
    all_cases = []

    # 1. 尝试下载并处理 CAIL 数据集
    for name, url in CAIL_URLS.items():
        dest = OUTPUT_DIR / f"{name}.json"
        if download_file(url, dest):
            try:
                cases = process_cail_json(dest)
                all_cases.extend(cases)
            except Exception as e:
                logger.error(f"处理 {name} 失败：{e}")

    if all_cases:
        df_real = pd.DataFrame(all_cases)
        real_path = OUTPUT_DIR / "cail_labor_cases.csv"
        df_real.to_csv(real_path, index=False, encoding="utf-8-sig")
        logger.info(f"真实案件数据保存至：{real_path}")
        logger.info(f"  男性：{(df_real['gender'] == '男').sum()} 条")
        logger.info(f"  女性：{(df_real['gender'] == '女').sum()} 条")
    else:
        logger.warning("未能从 CAIL 获取数据，将使用合成补充数据")

    # 2. 生成合成补充数据（保证实验有足够样本）
    df_synthetic = build_synthetic_supplement(n_per_gender=150)
    synth_path = OUTPUT_DIR / "synthetic_cases.csv"
    df_synthetic.to_csv(synth_path, index=False, encoding="utf-8-sig")
    logger.info(f"合成数据保存至：{synth_path}")

    # 3. 合并数据集
    dfs = []
    if all_cases:
        dfs.append(df_real)
    dfs.append(df_synthetic)
    df_combined = pd.concat(dfs, ignore_index=True)
    combined_path = OUTPUT_DIR / "combined_labor_cases.csv"
    df_combined.to_csv(combined_path, index=False, encoding="utf-8-sig")
    logger.info(f"合并数据集保存至：{combined_path}，共 {len(df_combined)} 条")


if __name__ == "__main__":
    main()
