import os
import sys
import json
import re
import argparse
import math
import inspect
from typing import Dict, List, Optional, Tuple

# ========== 单卡锁定（在导入 torch/transformers 前执行） ==========
def _extract_gpu_arg(argv: List[str], default: str = "0") -> str:
    for i, arg in enumerate(argv):
        if arg.startswith("--gpu="):
            return arg.split("=", 1)[1]
        if arg == "--gpu" and i + 1 < len(argv):
            return argv[i + 1]
    return default

env_vis = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
try:
    gpu_to_use = _extract_gpu_arg(sys.argv, default="0")
except Exception:
    gpu_to_use = "0"
# 若未设置或暴露了多卡，则强制只暴露单卡（默认0）以确保直接运行稳定
if (not env_vis) or ("," in env_vis):
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_to_use
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")

# 清理可能由外部启动器注入的分布式环境变量，避免误触多卡/分布式
for _k in ["RANK", "LOCAL_RANK", "WORLD_SIZE"]:
    os.environ.pop(_k, None)

import numpy as np
import torch
from torch.utils.data import Dataset
from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
import pandas as pd

from transformers import (
    AutoTokenizer,
    AutoModel,
    AutoModelForSequenceClassification,
    AutoConfig,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
    set_seed,
)
try:
    from transformers import EarlyStoppingCallback  # type: ignore
except Exception:  # pragma: no cover
    EarlyStoppingCallback = None  # type: ignore


def preprocess_text(text: str) -> str:
    if text is None:
        return ""
    text = str(text)
    text = re.sub(r"\{%.+?%\}", " ", text)
    text = re.sub(r"@.+?( |$)", " ", text)
    text = re.sub(r"【.+?】", " ", text)
    text = re.sub(r"\u200b", " ", text)
    text = re.sub(
        r"[\U0001F600-\U0001F64F\U0001F300-\U0001F5FF\U0001F680-\U0001F6FF\U0001F1E0-\U0001F1FF\U00002600-\U000027BF\U0001f900-\U0001f9ff\U0001f018-\U0001f270\U0000231a-\U0000231b\U0000238d-\U0000238d\U000024c2-\U0001f251]+",
        "",
        text,
    )
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def ensure_base_model_local(model_name_or_path: str, local_model_root: str) -> Tuple[str, AutoTokenizer]:
    os.makedirs(local_model_root, exist_ok=True)
    base_dir = os.path.join(local_model_root, "bert-base-chinese")

    def is_ready(path: str) -> bool:
        return os.path.isdir(path) and os.path.isfile(os.path.join(path, "config.json"))

    # 1) 本地现成
    if is_ready(base_dir):
        tokenizer = AutoTokenizer.from_pretrained(base_dir)
        return base_dir, tokenizer

    # 2) 本机缓存
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, local_files_only=True)
        base = AutoModel.from_pretrained(model_name_or_path, local_files_only=True)
        os.makedirs(base_dir, exist_ok=True)
        tokenizer.save_pretrained(base_dir)
        base.save_pretrained(base_dir)
        return base_dir, tokenizer
    except Exception:
        pass

    # 3) 远程下载
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    base = AutoModel.from_pretrained(model_name_or_path)
    os.makedirs(base_dir, exist_ok=True)
    tokenizer.save_pretrained(base_dir)
    base.save_pretrained(base_dir)
    return base_dir, tokenizer


class TextClassificationDataset(Dataset):
    def __init__(
        self,
        dataframe: pd.DataFrame,
        tokenizer: AutoTokenizer,
        text_column: str,
        label_column: str,
        label2id: Dict[str, int],
        max_length: int,
    ) -> None:
        self.dataframe = dataframe.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.text_column = text_column
        self.label_column = label_column
        self.label2id = label2id
        self.max_length = max_length

    def __len__(self) -> int:
        return len(self.dataframe)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        row = self.dataframe.iloc[idx]
        text = preprocess_text(row[self.text_column])
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,
            truncation=True,
            padding=False,
            return_tensors="pt",
        )
        item = {k: v.squeeze(0) for k, v in encoding.items()}
        if self.label_column in row and pd.notna(row[self.label_column]):
            label_str = str(row[self.label_column])
            item["labels"] = torch.tensor(self.label2id[label_str], dtype=torch.long)
        return item


def build_label_mappings(train_df: pd.DataFrame, label_column: str) -> Tuple[Dict[str, int], Dict[int, str]]:
    labels: List[str] = [str(x) for x in train_df[label_column].dropna().astype(str).tolist()]
    unique_sorted = sorted(set(labels))
    label2id = {label: i for i, label in enumerate(unique_sorted)}
    id2label = {i: label for label, i in label2id.items()}
    return label2id, id2label


def compute_metrics_fn(eval_pred) -> Dict[str, float]:
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="weighted", zero_division=0)
    acc = accuracy_score(labels, preds)
    return {
        "accuracy": float(acc),
        "precision": float(precision),
        "recall": float(recall),
        "f1": float(f1),
    }


def autodetect_columns(df: pd.DataFrame, text_col: str, label_col: str) -> Tuple[str, str]:
    if text_col != "auto" and label_col != "auto":
        return text_col, label_col
    candidates_text = ["text", "content", "sentence", "title", "desc", "question"]
    candidates_label = ["label", "labels", "category", "topic", "class"]
    t = text_col
    l = label_col
    if text_col == "auto":
        for name in candidates_text:
            if name in df.columns:
                t = name
                break
    if label_col == "auto":
        for name in candidates_label:
            if name in df.columns:
                l = name
                break
    if t == "auto" or l == "auto":
        raise ValueError(
            f"无法自动识别列名，请显式传入 --text_col 与 --label_col。现有列: {list(df.columns)}"
        )
    return t, l


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="使用 google-bert/bert-base-chinese 在本目录数据集上进行文本分类微调")
    parser.add_argument("--train_file", type=str, default="./dataset/web_text_zh_train.csv")
    parser.add_argument("--valid_file", type=str, default="./dataset/web_text_zh_valid.csv")
    parser.add_argument("--text_col", type=str, default="auto", help="文本列名，默认自动识别")
    parser.add_argument("--label_col", type=str, default="auto", help="标签列名，默认自动识别")
    parser.add_argument("--model_root", type=str, default="./model", help="本地模型根目录")
    parser.add_argument("--pretrained_name", type=str, default="google-bert/bert-base-chinese")
    parser.add_argument("--save_subdir", type=str, default="bert-chinese-classifier")
    parser.add_argument("--max_length", type=int, default=128)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--num_epochs", type=int, default=10)
    parser.add_argument("--learning_rate", type=float, default=2e-5)
    parser.add_argument("--weight_decay", type=float, default=0.01)
    parser.add_argument("--warmup_ratio", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--gpu", type=str, default=os.environ.get("CUDA_VISIBLE_DEVICES", "0"), help="指定单卡 GPU，如 0 或 1")
    parser.add_argument("--eval_fraction", type=float, default=0.25, help="每多少个 epoch 做一次评估与保存，例如 0.25 表示每四分之一个 epoch")
    parser.add_argument("--early_stop_patience", type=int, default=5, help="早停耐心（以评估轮次计）")
    parser.add_argument("--early_stop_threshold", type=float, default=0.0, help="早停最小改善阈值（与 metric_for_best_model 同单位）")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    set_seed(args.seed)

    script_dir = os.path.dirname(os.path.abspath(__file__))
    model_root = args.model_root if os.path.isabs(args.model_root) else os.path.join(script_dir, args.model_root)
    os.makedirs(model_root, exist_ok=True)

    # 确保基础模型就绪
    base_dir, tokenizer = ensure_base_model_local(args.pretrained_name, model_root)
    print(f"[Info] 使用基础模型目录: {base_dir}")

    # 读取数据
    train_path = args.train_file if os.path.isabs(args.train_file) else os.path.join(script_dir, args.train_file)
    valid_path = args.valid_file if os.path.isabs(args.valid_file) else os.path.join(script_dir, args.valid_file)
    if not os.path.isfile(train_path):
        raise FileNotFoundError(f"训练集不存在: {train_path}")
    train_df = pd.read_csv(train_path)
    if not os.path.isfile(valid_path):
        # 若未提供或不存在验证集，自动切分
        shuffled = train_df.sample(frac=1.0, random_state=args.seed).reset_index(drop=True)
        split_idx = int(len(shuffled) * 0.9)
        valid_df = shuffled.iloc[split_idx:].reset_index(drop=True)
        train_df = shuffled.iloc[:split_idx].reset_index(drop=True)
    else:
        valid_df = pd.read_csv(valid_path)
    print(f"[Info] 训练集: {train_path} | 样本数: {len(train_df)}")
    print(f"[Info] 验证集: {valid_path if os.path.isfile(valid_path) else '(从训练集切分)'} | 样本数: {len(valid_df)}")

    # 自动识别列名
    text_col, label_col = autodetect_columns(train_df, args.text_col, args.label_col)
    print(f"[Info] 文本列: {text_col} | 标签列: {label_col}")

    # 标签映射
    label2id, id2label = build_label_mappings(train_df, label_col)
    if len(label2id) < 2:
        raise ValueError("标签类别数少于 2，无法训练分类模型。")
    print(f"[Info] 标签类别数: {len(label2id)}")

    # 数据集
    train_dataset = TextClassificationDataset(train_df, tokenizer, text_col, label_col, label2id, args.max_length)
    eval_dataset = TextClassificationDataset(valid_df, tokenizer, text_col, label_col, label2id, args.max_length)
    collator = DataCollatorWithPadding(tokenizer=tokenizer)

    # 模型
    config = AutoConfig.from_pretrained(
        base_dir,
        num_labels=len(label2id),
        id2label={int(i): str(l) for i, l in id2label.items()},
        label2id={str(l): int(i) for l, i in label2id.items()},
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        base_dir,
        config=config,
        ignore_mismatched_sizes=True,
    )

    # 训练参数
    output_dir = os.path.join(model_root, args.save_subdir)
    os.makedirs(output_dir, exist_ok=True)
    # 训练参数（兼容不同 transformers 版本）
    args_dict = {
        "output_dir": output_dir,
        "per_device_train_batch_size": args.batch_size,
        "per_device_eval_batch_size": args.batch_size,
        "learning_rate": args.learning_rate,
        "weight_decay": args.weight_decay,
        "num_train_epochs": args.num_epochs,
        "logging_steps": 100,
        "fp16": args.fp16,
        "seed": args.seed,
    }

    sig = inspect.signature(TrainingArguments.__init__)
    allowed = set(sig.parameters.keys())

    # 可选参数（仅在支持时添加，尽量简化与参考实现一致以提升兼容性）
    if "warmup_ratio" in allowed:
        args_dict["warmup_ratio"] = args.warmup_ratio
    if "report_to" in allowed:
        args_dict["report_to"] = []
    # 评估/保存步进：按 eval_fraction 折算每个 epoch 的步数
    steps_per_epoch = max(1, math.ceil(len(train_dataset) / max(1, args.batch_size)))
    eval_every_steps = max(1, math.ceil(steps_per_epoch * max(0.01, min(1.0, args.eval_fraction))))
    # 策略式（新/旧版本字段名兼容）
    key_eval = "evaluation_strategy" if "evaluation_strategy" in allowed else ("eval_strategy" if "eval_strategy" in allowed else None)
    if key_eval:
        args_dict[key_eval] = "steps"
    if "save_strategy" in allowed:
        args_dict["save_strategy"] = "steps"
    if "eval_steps" in allowed:
        args_dict["eval_steps"] = eval_every_steps
    if "save_steps" in allowed:
        args_dict["save_steps"] = eval_every_steps
    if "save_total_limit" in allowed:
        args_dict["save_total_limit"] = 5
    # 将日志步长与评估/保存步长对齐，减少刷屏
    if "logging_steps" in allowed:
        args_dict["logging_steps"] = eval_every_steps
    # 最优模型回滚（仅当评估与保存策略一致时开启）
    if "metric_for_best_model" in allowed:
        args_dict["metric_for_best_model"] = "f1"
    if "greater_is_better" in allowed:
        args_dict["greater_is_better"] = True
    if "load_best_model_at_end" in allowed:
        eval_strat = args_dict.get("evaluation_strategy", args_dict.get("eval_strategy"))
        save_strat = args_dict.get("save_strategy")
        if eval_strat == save_strat and eval_strat in ("steps", "epoch"):
            args_dict["load_best_model_at_end"] = True

    # 兼容无 warmup_ratio 的版本：若支持 warmup_steps 则忽略比例
    if "warmup_ratio" not in allowed and "warmup_steps" in allowed:
        # 不计算总步数，默认 0
        args_dict["warmup_steps"] = 0

    # 若不支持策略式参数：退化为每 eval_every_steps 步保存/评估
    if "save_strategy" not in allowed and "save_steps" in allowed:
        args_dict["save_steps"] = eval_every_steps
    if ("evaluation_strategy" not in allowed and "eval_strategy" not in allowed) and "eval_steps" in allowed:
        args_dict["eval_steps"] = eval_every_steps

    # 如果支持 load_best_model_at_end，但无法同时设置评估/保存策略，则关闭它以避免报错
    if "load_best_model_at_end" in allowed:
        want_load_best = args_dict.get("load_best_model_at_end", False)
        eval_set = args_dict.get("evaluation_strategy", None)
        save_set = args_dict.get("save_strategy", None)
        if want_load_best and (eval_set is None or save_set is None or eval_set != save_set):
            args_dict["load_best_model_at_end"] = False

    training_args = TrainingArguments(**args_dict)
    print("[Info] 训练参数要点:")
    print(f"       epochs={args.num_epochs}, batch_size={args.batch_size}, lr={args.learning_rate}, weight_decay={args.weight_decay}")
    print(f"       max_length={args.max_length}, seed={args.seed}, fp16={args.fp16}")
    if "warmup_ratio" in allowed and "warmup_ratio" in args_dict:
        print(f"       warmup_ratio={args_dict['warmup_ratio']}")
    elif "warmup_steps" in allowed and "warmup_steps" in args_dict:
        print(f"       warmup_steps={args_dict['warmup_steps']}")
    print(f"       steps_per_epoch={steps_per_epoch}, eval_every_steps={eval_every_steps}")
    print(f"       eval_strategy={args_dict.get('evaluation_strategy', args_dict.get('eval_strategy'))}, save_strategy={args_dict.get('save_strategy')}, logging_steps={args_dict.get('logging_steps')}")
    print(f"       save_total_limit={args_dict.get('save_total_limit', 'n/a')}, load_best_model_at_end={args_dict.get('load_best_model_at_end', False)}")

    callbacks = []
    if EarlyStoppingCallback is not None and (args_dict.get("evaluation_strategy") in ("steps", "epoch") or "eval_steps" in allowed):
        try:
            callbacks.append(
                EarlyStoppingCallback(
                    early_stopping_patience=args.early_stop_patience,
                    early_stopping_threshold=args.early_stop_threshold,
                )
            )
        except Exception:
            pass

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        data_collator=collator,
        compute_metrics=compute_metrics_fn,
        callbacks=callbacks,
    )
    # 设备与 GPU 信息
    try:
        device_cnt = torch.cuda.device_count()
        dev_name = torch.cuda.get_device_name(0) if device_cnt > 0 else "cpu"
        print(f"[Info] CUDA 可见设备数: {device_cnt}, 当前设备: {dev_name}, CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}")
    except Exception:
        pass

    print("[Info] 开始训练 ...")

    trainer.train()

    # 保存
    tokenizer.save_pretrained(output_dir)
    trainer.model.config.id2label = {int(i): str(l) for i, l in id2label.items()}
    trainer.model.config.label2id = {str(l): int(i) for l, i in label2id.items()}
    trainer.save_model(output_dir)
    try:
        best_metric = getattr(trainer.state, "best_metric", None)
        best_ckpt = getattr(trainer.state, "best_model_checkpoint", None)
        if best_metric is not None and best_ckpt is not None:
            print(f"[Info] 最优模型: metric={best_metric:.6f} | checkpoint={best_ckpt}")
    except Exception:
        pass

    with open(os.path.join(output_dir, "label_map.json"), "w", encoding="utf-8") as f:
        json.dump(
            {"label2id": trainer.model.config.label2id, "id2label": trainer.model.config.id2label},
            f,
            ensure_ascii=False,
            indent=2,
        )

    # 训练曲线：可选保存训练与评估 loss
    try:
        import matplotlib.pyplot as plt  # type: ignore
        logs = trainer.state.log_history
        t_steps, t_losses, e_steps, e_losses = [], [], [], []
        step_counter = 0
        for rec in logs:
            if "loss" in rec and "epoch" in rec:
                step_counter += 1
                t_steps.append(step_counter)
                t_losses.append(rec["loss"])
            if "eval_loss" in rec:
                e_steps.append(step_counter)
                e_losses.append(rec["eval_loss"])
        if t_losses or e_losses:
            plt.figure(figsize=(8,4))
            if t_losses:
                plt.plot(t_steps, t_losses, label="train_loss")
            if e_losses:
                plt.plot(e_steps, e_losses, label="eval_loss")
            plt.xlabel("training step (logged)")
            plt.ylabel("loss")
            plt.legend()
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, "training_curve.png"))
    except Exception:
        pass

    print(f"微调完成，模型已保存到: {output_dir}")


if __name__ == "__main__":
    main()