from typing import Any


def build_lora_config(params: dict[str, Any]):
    """返回实际的 peft.LoraConfig 对象。"""
    from peft import LoraConfig, TaskType

    target_modules = params.get("lora_target_modules", "all-linear")
    if isinstance(target_modules, str):
        if target_modules == "all-linear":
            target_modules = ["linear", "lm_head", "q_proj", "v_proj", "k_proj", "o_proj"]

    return LoraConfig(
        r=params.get("lora_r", 16),
        lora_alpha=params.get("lora_alpha", 32),
        lora_dropout=params.get("lora_dropout", 0.05),
        target_modules=target_modules,
        task_type=TaskType.CAUSAL_LM,
    )


def build_qlora_config(params: dict[str, Any]):
    """返回 peft.LoraConfig 对象（量化已在 load_model 中通过 HQQ 处理）。"""
    from peft import LoraConfig, TaskType

    target_modules = params.get("lora_target_modules", "all-linear")
    if isinstance(target_modules, str) and target_modules == "all-linear":
        target_modules = ["linear", "lm_head", "q_proj", "v_proj", "k_proj", "o_proj"]

    return LoraConfig(
        r=params.get("lora_r", 16),
        lora_alpha=params.get("lora_alpha", 32),
        lora_dropout=params.get("lora_dropout", 0.05),
        target_modules=target_modules,
        task_type=TaskType.CAUSAL_LM,
    )


def build_adalora_config(params: dict[str, Any]):
    """返回实际的 peft.AdaLoraConfig 对象。"""
    from peft import AdaLoraConfig, TaskType

    target_modules = params.get("lora_target_modules", "all-linear")
    if isinstance(target_modules, str):
        if target_modules == "all-linear":
            target_modules = ["linear", "lm_head", "q_proj", "v_proj", "k_proj", "o_proj"]

    # total_step 必须由外部传入，AdaLoraConfig 的 __post_init__ 会校验 > 0
    # 如果没有传入，给一个较大的默认值（10000），train() 中会重新覆盖
    total_step = params.get("total_step", 10000)

    return AdaLoraConfig(
        init_r=params.get("adalora_init_r", 8),
        target_r=params.get("adalora_target_r", 16),
        beta1=params.get("adalora_beta1", 0.85),
        beta2=params.get("adalora_beta2", 0.85),
        target_modules=target_modules,
        task_type=TaskType.CAUSAL_LM,
        total_step=total_step,
    )