"""
DocuMint Smart Training Pipeline
- Core adapter (one-time training)
- Skill-wise adapters (additive learning)
- Safe continual learning (no destruction)
"""

import os
import gc
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)
from peft import (
    LoraConfig,
    get_peft_model,
    PeftModel,
    TaskType,
)
from huggingface_hub import login


# ================== CONFIG ==================

BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"

CORE_REPO = "himu1780/DocuMint-Core"
SKILL_REPO_PREFIX = "himu1780/DocuMint-Skill"

OUTPUT_DIR = "./lora_output"

MAX_LENGTH = 512
GRAD_ACCUM = 4
LOGGING_STEPS = 50
SAVE_STEPS = 500

TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]


# ================== UTILS ==================

def cleanup():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def hf_auth():
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise RuntimeError("HF_TOKEN not set")
    login(token=token)


# ================== DATA ==================

def format_example(ex):
    if "instruction" in ex and "output" in ex:
        text = (
            "<|im_start|>user\n"
            + ex["instruction"]
            + "<|im_end|>\n<|im_start|>assistant\n"
            + ex["output"]
            + "<|im_end|>"
        )
    elif "question" in ex and "answer" in ex:
        text = (
            "<|im_start|>user\n"
            + ex["question"]
            + "<|im_end|>\n<|im_start|>assistant\n"
            + ex["answer"]
            + "<|im_end|>"
        )
    else:
        text = ex.get("text", str(ex))

    return {"text": text}


def prepare_dataset(tokenizer, dataset_name):
    """
    Supports:
    - gsm8k
    - gsm8k:main
    - any_dataset
    """

    # Auto-fix gsm8k without config
    if dataset_name == "gsm8k":
        dataset_name = "gsm8k:main"

    # Handle dataset:config format
    if ":" in dataset_name:
        name, config = dataset_name.split(":", 1)
        dataset = load_dataset(name, config, split="train")
    else:
        dataset = load_dataset(dataset_name, split="train")

    dataset = dataset.map(format_example, remove_columns=dataset.column_names)

    def tokenize(ex):
        tokens = tokenizer(
            ex["text"],
            truncation=True,
            padding="max_length",
            max_length=MAX_LENGTH,
        )
        tokens["labels"] = tokens["input_ids"].copy()
        return tokens

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


# ================== MODEL ==================

def load_base():
    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL, trust_remote_code=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float32,   # CPU safe
        device_map="cpu",
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    return model, tokenizer


def lora_config():
    return LoraConfig(
        r=8,
        lora_alpha=16,
        lora_dropout=0.05,
        target_modules=TARGET_MODULES,
        task_type=TaskType.CAUSAL_LM,
        bias="none",
    )


# ================== ADAPTER LOGIC ==================

def load_core_adapter(model):
    core_path = os.path.join(OUTPUT_DIR, "core")

    if not os.path.exists(core_path):
        raise RuntimeError("Core adapter not found. Train core first.")

    model = PeftModel.from_pretrained(model, core_path)

    # Freeze everything
    for p in model.parameters():
        p.requires_grad = False

    print("🧠 Core adapter loaded and frozen")
    return model


def load_or_create_adapter(model, skill_name):
    adapter_path = os.path.join(OUTPUT_DIR, skill_name)

    if os.path.exists(adapter_path):
        print(f"🔁 Loading existing adapter: {skill_name}")
        model = PeftModel.from_pretrained(
            model, adapter_path, is_trainable=True
        )
    else:
        print(f"🆕 Creating new adapter: {skill_name}")
        model = get_peft_model(model, lora_config())

    model.print_trainable_parameters()
    return model


# ================== TRAIN ==================

def train_skill(
    dataset_name: str,
    skill_name: str,
    epochs: int,
    lr: float,
    batch_size: int,
):
    """
    skill_name:
    - "core"  -> core training (one time)
    - others  -> skill training (requires core)
    """

    hf_auth()

    model, tokenizer = load_base()

    # IMPORTANT FIX:
    # Load core ONLY if training a skill
    if skill_name != "core":
        model = load_core_adapter(model)

    # Load or create adapter
    model = load_or_create_adapter(model, skill_name)

    dataset = prepare_dataset(tokenizer, dataset_name)

    args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=GRAD_ACCUM,
        learning_rate=lr,
        logging_steps=LOGGING_STEPS,
        save_steps=SAVE_STEPS,
        save_total_limit=2,
        fp16=False,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        report_to="none",
        remove_unused_columns=False,
    )

    collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=dataset,
        data_collator=collator,
    )

    trainer.train()

    # Save locally
    save_path = os.path.join(OUTPUT_DIR, skill_name)
    model.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)

    # Push to Hub
    if skill_name == "core":
        repo = CORE_REPO
    else:
        repo = f"{SKILL_REPO_PREFIX}-{skill_name}"

    model.push_to_hub(repo)
    tokenizer.push_to_hub(repo)

    cleanup()
    print(f"✅ Training finished for adapter: {skill_name}")


# ================== ROUTING (INFERENCE) ==================

def load_for_inference(skill_name: str):
    model, tokenizer = load_base()

    model = PeftModel.from_pretrained(model, CORE_REPO)
    model = PeftModel.from_pretrained(
        model, f"{SKILL_REPO_PREFIX}-{skill_name}"
    )

    model.eval()
    print(f"🚦 Routed adapters: Core + {skill_name}")
    return model, tokenizer


# ================== MAIN ==================

if __name__ == "__main__":
    print("🏆 DocuMint Smart Training System Ready")
    print("Use train_skill() to train core or add skills safely")