"""Fine-tuning entry point for CORTEX AI.

This script is written to run on GPU hardware. It is NOT executed in this
repository: the sandbox has 4 CPU cores, ~15 GiB of RAM and no accelerator, and
the checkpoint is 893 GB. Treat it as the documented path, not a verified run.

Usage::

    python -m cortex_ai.training.train --config train_config.json
    python -m cortex_ai.training.train --dataset data/cortex_sft.jsonl --epochs 1

Requires: torch, transformers, peft, datasets
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from .config import TrainConfig
from .dataset import dataset_stats, iter_encoded, load_jsonl


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Affinage de CORTEX AI")
    parser.add_argument("--config", type=str, default="", help="fichier JSON de configuration")
    parser.add_argument("--dataset", type=str, default="", help="dataset JSONL")
    parser.add_argument("--output-dir", type=str, default="", help="dossier de sortie")
    parser.add_argument("--epochs", type=int, default=0)
    parser.add_argument("--dry-run", action="store_true", help="valider les donnees sans entrainer")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_arg_parser().parse_args(argv)

    config = TrainConfig.load(args.config) if args.config else TrainConfig()
    if args.dataset:
        config.dataset_path = args.dataset
    if args.output_dir:
        config.output_dir = args.output_dir
    if args.epochs:
        config.epochs = args.epochs

    examples = load_jsonl(config.dataset_path)
    stats = dataset_stats(examples)
    print("=== Donnees ===")
    for key, value in stats.items():
        print(f"  {key:20s} {value}")

    # Encoding every example catches format problems before a GPU is booked.
    encoded = list(iter_encoded(examples, config.thinking_mode))
    print(f"  {'exemples encodes':20s} {len(encoded)}")
    print(f"  {'prompt le plus long':20s} {max(len(e) for e in encoded)} caracteres")

    if args.dry_run:
        print("\n--dry-run : validation terminee, aucun entrainement lance.")
        return 0

    try:
        import torch
        from datasets import Dataset
        from peft import LoraConfig, get_peft_model
        from transformers import (
            AutoModelForCausalLM,
            AutoTokenizer,
            DataCollatorForLanguageModeling,
            Trainer,
            TrainingArguments,
        )
    except ImportError as exc:
        print(
            f"\nDependance manquante : {exc.name}.\n"
            "Installez torch, transformers, peft et datasets, puis relancez.\n"
            "Sur cette machine (sans GPU, 893 Go de checkpoint), l'entrainement "
            "n'est pas executable : utilisez --dry-run pour valider les donnees.",
            file=sys.stderr,
        )
        return 2

    if not torch.cuda.is_available():
        print(
            "\nAucun GPU detecte. L'affinage de ce checkpoint exige plusieurs "
            "accelerateurs. Utilisez --dry-run pour valider les donnees.",
            file=sys.stderr,
        )
        return 2

    print("\n=== Chargement du modele ===")
    tokenizer = AutoTokenizer.from_pretrained(config.model_id)
    model = AutoModelForCausalLM.from_pretrained(
        config.model_id, torch_dtype=torch.bfloat16, device_map="auto"
    )

    if config.method == "lora":
        model = get_peft_model(
            model,
            LoraConfig(
                r=config.lora.r,
                lora_alpha=config.lora.alpha,
                lora_dropout=config.lora.dropout,
                target_modules=config.lora.target_modules,
                task_type="CAUSAL_LM",
            ),
        )
        model.print_trainable_parameters()

    dataset = Dataset.from_dict({"text": encoded})

    def tokenize(batch):
        return tokenizer(batch["text"], truncation=True, max_length=config.max_seq_length)

    dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])

    training_args = TrainingArguments(
        output_dir=config.output_dir,
        num_train_epochs=config.epochs,
        per_device_train_batch_size=config.batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        warmup_ratio=config.warmup_ratio,
        weight_decay=config.weight_decay,
        logging_steps=config.logging_steps,
        save_steps=config.save_steps,
        bf16=True,
        seed=config.seed,
        report_to=[],
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )
    trainer.train()
    trainer.save_model(config.output_dir)
    tokenizer.save_pretrained(config.output_dir)
    print(f"\nModele enregistre dans {config.output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())