# core/finetune_engine.py
"""
Fine-Tuning Engine – Utilise Unsloth pour entraîner un LoRA sur le dataset.
Version finale avec SFTConfig et désactivation des sauvegardes intermédiaires.
"""
import os
import logging
import subprocess
import json
from pathlib import Path
from typing import Optional

logger = logging.getLogger("lucie.finetune_engine")

class FinetuneEngine:
    def __init__(self, model_name: str = "Qwen/Qwen2.5-0.5B-Instruct",
                 data_dir: str = "/data/datasets",
                 output_dir: str = "/data/lora"):
        self.model_name = model_name
        self.data_dir = Path(data_dir)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def prepare_config(self, dataset_path: Path) -> Path:
        config = {
            "model_name": self.model_name,
            "dataset": str(dataset_path),
            "output_dir": str(self.output_dir),
            "lora_r": 16,
            "lora_alpha": 32,
            "lora_dropout": 0.0,
            "learning_rate": 2e-4,
            "num_train_epochs": 1,
            "per_device_train_batch_size": 4,
            "gradient_accumulation_steps": 4,
            "save_steps": 1000000,
            "logging_steps": 1,
            "fp16": True,
            "max_seq_length": 512,
        }
        config_path = self.data_dir / "finetune_config.json"
        with open(config_path, "w") as f:
            json.dump(config, f, indent=2)
        return config_path

    def run_unsloth(self, config_path: Path) -> bool:
        script = f"""
import unsloth
from unsloth import FastLanguageModel
from trl import SFTConfig, SFTTrainer
import json
import torch

with open("{config_path}") as f:
    config = json.load(f)

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=config["model_name"],
    max_seq_length=config["max_seq_length"],
    load_in_4bit=True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=config["lora_r"],
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_alpha=config["lora_alpha"],
    lora_dropout=config["lora_dropout"],
)

from datasets import load_dataset
dataset = load_dataset("json", data_files=config["dataset"], split="train")

def formatting_func(examples):
    instructions = examples["instruction"]
    responses = examples["response"]
    texts = []
    for ins, resp in zip(instructions, responses):
        texts.append(f"### Instruction:\\n{{ins}}\\n### Response:\\n{{resp}}")
    return texts

training_args = SFTConfig(
    output_dir=config["output_dir"],
    per_device_train_batch_size=config["per_device_train_batch_size"],
    gradient_accumulation_steps=config["gradient_accumulation_steps"],
    num_train_epochs=config["num_train_epochs"],
    learning_rate=config["learning_rate"],
    fp16=config["fp16"],
    logging_steps=config["logging_steps"],
    save_steps=config["save_steps"],
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    formatting_func=formatting_func,
    max_seq_length=config["max_seq_length"],
    args=training_args,
)

trainer.train()
model.save_pretrained(config["output_dir"])
tokenizer.save_pretrained(config["output_dir"])
print("✅ Fine-tuning terminé.")
"""
        script_path = self.data_dir / "train.py"
        with open(script_path, "w") as f:
            f.write(script)

        try:
            result = subprocess.run(
                ["python", str(script_path)],
                capture_output=True,
                text=True,
                timeout=7200
            )
            if result.returncode == 0:
                logger.info("✅ Fine-tuning réussi.")
                return True
            else:
                logger.error(f"❌ Fine-tuning échoué : {result.stderr}")
                if result.stdout:
                    logger.info(f"Sortie : {result.stdout}")
                return False
        except Exception as e:
            logger.error(f"❌ Erreur fine-tuning : {e}")
            return False

    def run(self, dataset: Path) -> bool:
        logger.info(f"🚀 Lancement du fine-tuning sur {dataset}...")
        config_path = self.prepare_config(dataset)
        return self.run_unsloth(config_path)