"""
预训练模型微调模块 — Person B 负责实现

功能要求:
1. load_pretrained_model: 加载 NLLB / mBART 等预训练翻译模型
2. setup_lora: 配置 LoRA 参数高效微调

技术要点:
- 使用 HuggingFace transformers 加载预训练模型
- 使用 PEFT (Parameter-Efficient Fine-Tuning) 库配置 LoRA
- 支持 NLLB-200 (Meta, 200种语言) 和 mBART-50 (Meta, 50种语言)
- 冻结预训练参数，只训练 LoRA adapter

这是实验对比的关键部分:
- 从头训练 vs 预训练微调
- Full fine-tuning vs LoRA
"""

from __future__ import annotations

import logging
from typing import Optional

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


def load_pretrained_model(
    model_name: str = "facebook/nllb-200-distilled-600M",
    src_lang: str = "eng_Latn",
    tgt_lang: str = "zho_Hans",
    device: str = "cuda",
) -> tuple:
    """
    加载预训练翻译模型和分词器。

    TODO [Person B]: 实现以下逻辑:
    1. 使用 AutoModelForSeq2SeqLM.from_pretrained(model_name) 加载模型
    2. 使用 AutoTokenizer.from_pretrained(model_name) 加载分词器
    3. 设置源语言和目标语言
    4. 将模型移到指定设备

    支持的模型:
    - facebook/nllb-200-distilled-600M (推荐，轻量级)
    - facebook/nllb-200-1.3B (中等规模)
    - facebook/mbart-large-50-many-to-many-mmt

    Returns:
        (model, tokenizer)
    """
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    # 设置语言对
    if hasattr(tokenizer, "lang_code_to_id"):
        tokenizer.src_lang = src_lang
        tokenizer.tgt_lang = tgt_lang
    if hasattr(model.config, "forced_bos_token_id"):
        model.config.forced_bos_token_id = tokenizer.lang_code_to_id.get(
            tgt_lang, tokenizer.bos_token_id
        )

    model = model.to(device)
    logger.info(f"Loaded pretrained model: {model_name}")
    return model, tokenizer


def setup_lora(
    model: nn.Module,
    r: int = 16,
    alpha: int = 32,
    dropout: float = 0.05,
    target_modules: Optional[list[str]] = None,
) -> nn.Module:
    """
    为模型配置 LoRA 微调。

    Returns:
        peft_model: LoRA 包装后的模型
    """
    from peft import LoraConfig, get_peft_model, TaskType

    if target_modules is None:
        target_modules = ["q_proj", "v_proj"]

    config = LoraConfig(
        r=r,
        lora_alpha=alpha,
        lora_dropout=dropout,
        target_modules=target_modules,
        task_type=TaskType.SEQ_2_SEQ_LM,
    )

    model = get_peft_model(model, config)

    # 打印可训练参数信息
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    ratio = 100 * trainable_params / total_params if total_params > 0 else 0
    logger.info(
        f"LoRA setup: trainable params={trainable_params:,} "
        f"/ total={total_params:,} ({ratio:.4f}%)"
    )

    return model


def freeze_model_except_lora(model: nn.Module):
    """冻结模型所有参数，只保留 LoRA 参数可训练。"""
    for name, param in model.named_parameters():
        if "lora" not in name.lower():
            param.requires_grad = False