"""FFT parcial do DeepSeek-V4-Flash-0731: experts MoE congelados, FFT dos modulos densos.
Usa Unsloth (kernels + packing) + TRL SFTTrainer + DeepSpeed Ulysses SP=8 para 64k ctx."""
import argparse
import json
import os
import sys
import time
import math
import torch
from datasets import load_from_disk
from transformers import AutoTokenizer, TrainingArguments
from trl import SFTTrainer, SFTConfig

try:
    from unsloth import FastModel
    UNSLOTH_AVAILABLE = True
except Exception:
    UNSLOTH_AVAILABLE = False


def freeze_moe(model):
    frozen = 0
    trainable = 0
    for name, p in model.named_parameters():
        is_moe = any(k in name.lower() for k in ("moe", "experts", "gate", "router", "moe_layer"))
        if is_moe:
            p.requires_grad_(False)
            frozen += p.numel()
        else:
            p.requires_grad_(True)
            trainable += p.numel()
    return frozen, trainable


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="hf_model/config.json")
    ap.add_argument("--resume", default="auto", choices=["auto", "true", "false"])
    ap.add_argument("--max-steps", type=int, default=-1)
    args = ap.parse_args()

    with open(args.config) as fh:
        cfg = json.load(fh)

    model_id = cfg["model_id"]
    max_seq_len = cfg["max_seq_len"]
    data_dir = cfg["data_dir"]
    out_dir = cfg["output_dir"]

    print(f"[train] model={model_id} max_seq_len={max_seq_len}", flush=True)
    print(f"[train] unsloth available: {UNSLOTH_AVAILABLE}", flush=True)

    if UNSLOTH_AVAILABLE:
        model, tok = FastModel.from_pretrained(
            model_name=model_id,
            max_seq_length=max_seq_len,
            dtype=torch.bfloat16,
            load_in_4bit=False,
            trust_remote_code=True,
        )
        FastModel.for_training(model)
    else:
        tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        from transformers import AutoModelForCausalLM
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            attn_implementation="flash_attention_2",
        )

    frozen, trainable = freeze_moe(model)
    print(f"[train] frozen params (MoE): {frozen:,}", flush=True)
    print(f"[train] trainable params (non-MoE): {trainable:,}", flush=True)
    print(f"[train] trainable %: {100*trainable/(frozen+trainable):.3f}%", flush=True)

    ds = load_from_disk(data_dir)
    print(f"[train] dataset: {len(ds)} samples", flush=True)

    sft_config = SFTConfig(
        dataset_text_field="text",
        max_length=max_seq_len,
        packing=cfg.get("packing", True),
        dataset_num_proc=8,
        report_to="none",
        output_dir=out_dir,
        num_train_epochs=cfg.get("epochs", 1),
        per_device_train_batch_size=cfg.get("micro_batch", 1),
        gradient_accumulation_steps=cfg.get("grad_accum", 8),
        learning_rate=cfg.get("lr", 1e-5),
        warmup_steps=cfg.get("warmup_steps", 50),
        lr_scheduler_type="cosine",
        bf16=True,
        logging_steps=10,
        save_steps=cfg.get("save_steps", 500),
        save_total_limit=cfg.get("save_total_limit", 4),
        save_strategy="steps",
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        max_steps=args.max_steps if args.max_steps > 0 else -1,
        optim="adamw_torch_fused",
        weight_decay=cfg.get("weight_decay", 0.01),
        max_grad_norm=cfg.get("max_grad_norm", 1.0),
        seed=42,
        dataloader_num_workers=4,
        remove_unused_columns=True,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tok,
        train_dataset=ds,
        args=sft_config,
    )

    resume_from = None
    if args.resume in ("auto", "true"):
        if os.path.isdir(out_dir):
            ckpts = sorted(
                [d for d in os.listdir(out_dir) if d.startswith("checkpoint-")],
                key=lambda x: int(x.split("-")[1]),
            )
            if ckpts:
                resume_from = os.path.join(out_dir, ckpts[-1])
                print(f"[train] resuming from {resume_from}", flush=True)

    t0 = time.time()
    if resume_from:
        trainer.train(resume_from_checkpoint=resume_from)
    else:
        trainer.train()
    train_secs = time.time() - t0

    print(f"[train] training done in {train_secs:.1f}s", flush=True)

    if args.max_steps <= 0:
        print(f"[train] saving final model to {out_dir}/final", flush=True)
        trainer.save_model(os.path.join(out_dir, "final"))
        tok.save_pretrained(os.path.join(out_dir, "final"))

    log = trainer.state.log_history if hasattr(trainer.state, "log_history") else []
    summary = {
        "train_secs": train_secs,
        "global_step": trainer.state.global_step,
        "max_steps": trainer.state.max_steps,
        "log_history": log,
        "frozen_params": frozen,
        "trainable_params": trainable,
        "model_id": model_id,
        "max_seq_len": max_seq_len,
    }
    summary_path = os.path.join(out_dir, "training_summary.json")
    os.makedirs(out_dir, exist_ok=True)
    with open(summary_path, "w") as fh:
        json.dump(summary, fh, indent=2, default=str)
    print(f"[train] summary saved to {summary_path}", flush=True)


if __name__ == "__main__":
    main()