import os, sys, argparse, logging
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig, TaskType

if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.is_available()}, devices={torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        print(f"  Device {i}: {torch.cuda.get_device_name(i)}")

def to_messages(example):
    messages = []
    system = example.get("system", "")
    if system and str(system).strip():
        messages.append({"role": "system", "content": str(system).strip()})
    messages.append({"role": "user", "content": str(example["user"]).strip()})
    messages.append({"role": "assistant", "content": str(example["assistant"]).strip()})
    return {"messages": messages}

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", default="Qwen/Qwen2.5-7B-Instruct")
    parser.add_argument("--dataset_name", default="Josephgflowers/Finance-Instruct-500k")
    parser.add_argument("--output_dir", default="/tmp/output/finance-sft")
    parser.add_argument("--hub_model_id", default="shah-shazid-askary/amd-finance-llm")
    parser.add_argument("--learning_rate", type=float, default=1e-5)
    parser.add_argument("--num_train_epochs", type=int, default=3)
    parser.add_argument("--warmup_ratio", type=float, default=0.1)
    parser.add_argument("--max_seq_length", type=int, default=8192)
    parser.add_argument("--per_device_train_batch_size", type=int, default=1)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=16)
    parser.add_argument("--bf16", action="store_true", default=True)
    parser.add_argument("--use_lora", action="store_true", default=False)
    parser.add_argument("--max_samples", type=int, default=None)
    parser.add_argument("--logging_steps", type=int, default=10)
    parser.add_argument("--save_steps", type=int, default=500)
    parser.add_argument("--eval_steps", type=int, default=500)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--push_to_hub", action="store_true", default=True)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    logger = logging.getLogger(__name__)

    logger.info("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    logger.info("Loading dataset...")
    ds = load_dataset(args.dataset_name, split="train")
    if args.max_samples:
        ds = ds.select(range(min(args.max_samples, len(ds))))
    ds = ds.map(to_messages, remove_columns=ds.column_names, batched=False)
    ds = ds.train_test_split(test_size=0.05, seed=args.seed)
    logger.info(f"Train: {len(ds["train"])} | Eval: {len(ds["test"])}")

    logger.info("Loading model...")
    attn_impl = "flash_attention_2" if torch.cuda.is_available() else "eager"
    model_kwargs = {
        "torch_dtype": torch.bfloat16 if args.bf16 else torch.float32,
        "attn_implementation": attn_impl,
        "trust_remote_code": True,
    }
    if torch.cuda.device_count() >= 1:
        model_kwargs["device_map"] = "auto"
    model = AutoModelForCausalLM.from_pretrained(args.model_name, **model_kwargs)

    peft_config = None
    if args.use_lora:
        peft_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM, inference_mode=False,
            r=32, lora_alpha=16, lora_dropout=0.05,
            target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        )
        model.print_trainable_parameters()

    sft_args = SFTConfig(
        output_dir=args.output_dir, num_train_epochs=args.num_train_epochs,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate, warmup_ratio=args.warmup_ratio,
        lr_scheduler_type="cosine", max_seq_length=args.max_seq_length,
        bf16=args.bf16,
        logging_steps=args.logging_steps, logging_strategy="steps",
        logging_first_step=True, save_steps=args.save_steps,
        eval_strategy="steps", eval_steps=args.eval_steps,
        load_best_model_at_end=True, metric_for_best_model="eval_loss",
        greater_is_better=False, seed=args.seed,
        push_to_hub=args.push_to_hub, hub_model_id=args.hub_model_id,
        disable_tqdm=True, gradient_checkpointing=True,
        run_name=f"finance-sft",
    )

    trainer = SFTTrainer(
        model=model, tokenizer=tokenizer,
        train_dataset=ds["train"], eval_dataset=ds["test"],
        args=sft_args, peft_config=peft_config,
    )
    logger.info("Starting training...")
    trainer.train()
    logger.info("Saving model...")
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    if args.push_to_hub:
        trainer.push_to_hub()
    logger.info("Done.")

if __name__ == "__main__":
    main()