from __future__ import annotations

import json
from pathlib import Path

from .constants import (
    DEFAULT_ADAPTER_DIR,
    DEFAULT_PACK_DIR,
    DEFAULT_MAX_SEQ_LEN,
    GENESIS_DIR,
    LORA_ALPHA,
    LORA_RANK,
    SMOKE_MAX_SEQ_LEN,
)


def train_sft(
    *,
    pack_path: Path,
    output_dir: Path = DEFAULT_ADAPTER_DIR,
    model_dir: Path = GENESIS_DIR,
    max_seq_len: int = DEFAULT_MAX_SEQ_LEN,
    max_steps: int | None = None,
    num_epochs: float = 1.0,
    per_device_batch_size: int = 1,
    grad_accum: int = 8,
    lr: float = 1e-4,
    lora_rank: int = LORA_RANK,
    lora_alpha: int = LORA_ALPHA,
    smoke: bool = False,
) -> Path:
    """LoRA SFT on genesis. Assistant/completion tokens only."""
    from local_eval.cuda_env import apply as apply_cuda

    apply_cuda()
    pack_path = Path(pack_path)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    if smoke:
        max_seq_len = min(max_seq_len, SMOKE_MAX_SEQ_LEN)
        max_steps = max_steps or 20

    rows = _load_pack(pack_path)
    if not rows:
        raise ValueError(f"empty pack: {pack_path}")

    import torch
    from datasets import Dataset
    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from trl import SFTConfig, SFTTrainer

    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=False)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    dataset = Dataset.from_list(
        [{"prompt": row["prompt"], "completion": row["completion"]} for row in rows]
    )

    model = AutoModelForCausalLM.from_pretrained(
        str(model_dir),
        torch_dtype=torch.bfloat16,
        trust_remote_code=False,
        attn_implementation="sdpa",
    )
    model.config.use_cache = False
    if hasattr(model, "enable_input_require_grads"):
        model.enable_input_require_grads()
    targets = lora_target_modules(model)
    model = get_peft_model(
        model,
        LoraConfig(
            r=lora_rank,
            lora_alpha=lora_alpha,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=targets,
        ),
    )
    model.print_trainable_parameters()

    args_kwargs = dict(
        output_dir=str(output_dir),
        bf16=True,
        learning_rate=lr,
        per_device_train_batch_size=per_device_batch_size,
        gradient_accumulation_steps=grad_accum,
        gradient_checkpointing=True,
        logging_steps=1,
        save_steps=max(max_steps or 200, 50),
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        report_to=[],
        max_length=max_seq_len,
        packing=False,
        completion_only_loss=True,
        remove_unused_columns=False,
    )
    if max_steps:
        args_kwargs["max_steps"] = max_steps
    else:
        args_kwargs["num_train_epochs"] = num_epochs
    config = SFTConfig(**_filter_kwargs(SFTConfig, args_kwargs))

    trainer = SFTTrainer(
        model=model,
        args=config,
        train_dataset=dataset,
        processing_class=tokenizer,
    )
    trainer.train()
    trainer.save_model(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))
    (output_dir / "sft-report.json").write_text(
        json.dumps(
            {
                "pack": str(pack_path),
                "n": len(rows),
                "max_steps": max_steps,
                "max_seq_len": max_seq_len,
                "lora_rank": lora_rank,
                "target_modules": targets,
                "smoke": smoke,
            },
            indent=2,
        )
        + "\n"
    )
    print(f"adapter: {output_dir}", flush=True)
    return output_dir


def lora_target_modules(model) -> list[str]:
    import torch

    wanted = {
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
        "in_proj_qkv",
        "in_proj",
        "out_proj",
        "gate",
    }
    found: set[str] = set()
    for name, module in model.named_modules():
        if not isinstance(module, torch.nn.Linear):
            continue
        leaf = name.rsplit(".", 1)[-1]
        if leaf in wanted:
            found.add(leaf)
    if not found:
        found = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}
    return sorted(found)


def default_pack(pack_dir: Path = DEFAULT_PACK_DIR) -> Path:
    packs = sorted(Path(pack_dir).glob("sft-*.jsonl"), key=lambda p: p.stat().st_mtime)
    if not packs:
        raise FileNotFoundError(f"no sft-*.jsonl under {pack_dir}")
    return packs[-1]


def _load_pack(path: Path) -> list[dict]:
    rows = []
    for line in Path(path).read_text().splitlines():
        if line.strip():
            rows.append(json.loads(line))
    return rows


def _filter_kwargs(cls, kwargs: dict) -> dict:
    try:
        fields = set(cls.__dataclass_fields__)
    except Exception:
        return kwargs
    return {key: value for key, value in kwargs.items() if key in fields}