"""
train.py — MamunAI LoRA Fine-Tuning (Resumable + GitHub Dataset)
- Downloads dataset at runtime from GitHub (never stored in repo).
- Resumes from existing LoRA adapter on Hugging Face Hub if available.
- Never auto-merges; merge is a separate explicit command.
Owner: Al Mamun Khan
"""

import json
import os
import logging
import shutil
import urllib.request
import torch
from pathlib import Path
from typing import Optional, Callable

from huggingface_hub import login, create_repo, hf_hub_download, list_repo_files, HfApi
from huggingface_hub.utils import RepositoryNotFoundError, EntryNotFoundError

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("MamunAI-Train")

# ── Constants ─────────────────────────────────────────────────────────────────

BASE_MODEL_ID       = "Qwen/Qwen2.5-0.5B-Instruct"
HF_REPO             = "almamunkhan/MamunAI"
GITHUB_DATASET_URL  = "https://raw.githubusercontent.com/hunterking42/Mamun_AI/main/dataset.jsonl"
RUNTIME_DATASET_DIR = "./runtime_datasets"
RUNTIME_DATASET     = f"{RUNTIME_DATASET_DIR}/dataset.jsonl"
OUTPUT_DIR          = "./lora_adapter"
MERGED_DIR          = "./merged_model"

MAX_SEQ_LENGTH = 256
LORA_RANK      = 8
LORA_ALPHA     = 16
LORA_DROPOUT   = 0.05
LEARNING_RATE  = 3e-4
NUM_EPOCHS     = 1
BATCH_SIZE     = 8
GRAD_ACCUM     = 4
WARMUP_STEPS   = 5
SAVE_STEPS     = 100

LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]

# ── System prompt ─────────────────────────────────────────────────────────────

SYSTEM_PROMPT = (
    "You are MamunAI, a personal AI assistant created by Al Mamun Khan "
    "from Bangladesh. You are helpful, honest, and concise. "
    "You respond in English or Bengali based on the user's language."
)

# ── Hardware detection ────────────────────────────────────────────────────────

HAS_GPU   = torch.cuda.is_available()
HAS_BF16  = HAS_GPU and torch.cuda.is_bf16_supported()
USE_FP16  = HAS_GPU and not HAS_BF16
USE_BF16  = HAS_GPU and HAS_BF16

# ── Log helper ────────────────────────────────────────────────────────────────

def log(msg: str, cb: Optional[Callable] = None) -> None:
    logger.info(msg)
    if cb:
        cb(msg)

# ── Dataset download ──────────────────────────────────────────────────────────

def download_dataset(cb: Optional[Callable] = None) -> list[dict]:
    """
    Fetch dataset.jsonl from GitHub at runtime.
    Always overwrites any existing local copy.
    Raises RuntimeError (with a clear message) on any failure.
    """
    log("Downloading dataset from GitHub...", cb)
    os.makedirs(RUNTIME_DATASET_DIR, exist_ok=True)

    try:
        urllib.request.urlretrieve(GITHUB_DATASET_URL, RUNTIME_DATASET)
    except Exception as exc:
        raise RuntimeError(
            f"GitHub dataset download failed: {exc}\n"
            f"URL checked: {GITHUB_DATASET_URL}"
        ) from exc

    if not os.path.isfile(RUNTIME_DATASET) or os.path.getsize(RUNTIME_DATASET) == 0:
        raise RuntimeError(
            "Downloaded dataset.jsonl is empty. "
            "Check the GitHub URL and repository visibility."
        )

    records: list[dict] = []
    with open(RUNTIME_DATASET, "r", encoding="utf-8") as fh:
        for lineno, raw in enumerate(fh, 1):
            raw = raw.strip()
            if not raw:
                continue
            try:
                records.append(json.loads(raw))
            except json.JSONDecodeError as exc:
                raise RuntimeError(
                    f"dataset.jsonl parse error at line {lineno}: {exc}"
                ) from exc

    if not records:
        raise RuntimeError(
            "dataset.jsonl was downloaded but contains zero valid examples."
        )

    log("Dataset downloaded successfully.", cb)
    log(f"Dataset contains {len(records)} training examples.", cb)
    return records

# ── Format helper ─────────────────────────────────────────────────────────────

def format_example(record: dict) -> str:
    return (
        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{record.get('instruction', '')}<|im_end|>\n"
        f"<|im_start|>assistant\n{record.get('output', '')}<|im_end|>"
    )

# ── HF Hub helpers ────────────────────────────────────────────────────────────

def adapter_exists_on_hub(repo_id: str, token: str) -> bool:
    try:
        files = list(list_repo_files(repo_id, repo_type="model", token=token))
        return "adapter_config.json" in files
    except RepositoryNotFoundError:
        return False
    except Exception as exc:
        logger.warning("Could not check Hub for existing adapter: %s", exc)
        return False


def download_adapter_from_hub(repo_id: str, token: str, local_dir: str) -> bool:
    """Download all known adapter files; return True if adapter_config.json landed."""
    filenames = [
        "adapter_config.json",
        "adapter_model.bin",
        "adapter_model.safetensors",
        "tokenizer_config.json",
        "tokenizer.json",
        "tokenizer.model",
        "special_tokens_map.json",
        "training_meta.json",
    ]
    os.makedirs(local_dir, exist_ok=True)
    got_config = False
    for name in filenames:
        try:
            hf_hub_download(
                repo_id=repo_id, filename=name,
                repo_type="model", token=token,
                local_dir=local_dir,
            )
            logger.info("  ↓ %s", name)
            if name == "adapter_config.json":
                got_config = True
        except EntryNotFoundError:
            pass
        except Exception as exc:
            logger.warning("  Could not download %s: %s", name, exc)
    return got_config

# ── Model loading ─────────────────────────────────────────────────────────────

def load_base_model_and_tokenizer(cb: Optional[Callable] = None):
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

    log(f"Loading tokenizer: {BASE_MODEL_ID}", cb)
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
    tokenizer.padding_side = "right"

    bnb_config: Optional[BitsAndBytesConfig] = None
    device_map = "auto" if HAS_GPU else "cpu"

    if HAS_GPU:
        try:
            import bitsandbytes  # noqa: F401
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16 if HAS_BF16 else torch.float16,
            )
            log("QLoRA 4-bit quantization enabled (GPU)", cb)
        except ImportError:
            log("bitsandbytes not installed — loading in full precision", cb)
    else:
        log("CUDA not available — running in CPU mode (slower)", cb)

    log(f"Loading model: {BASE_MODEL_ID}", cb)
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        quantization_config=bnb_config,
        device_map=device_map,
        trust_remote_code=True,
        dtype=torch.bfloat16 if HAS_BF16 else torch.float16 if HAS_GPU else torch.float32,
        low_cpu_mem_usage=True,
        
    )
    return model, tokenizer


def _prep_kbit(model):
    from peft import prepare_model_for_kbit_training
    try:
        return prepare_model_for_kbit_training(model)
    except Exception:
        return model


def apply_lora_fresh(model, cb: Optional[Callable] = None):
    from peft import LoraConfig, get_peft_model, TaskType
    model = _prep_kbit(model)
    config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=LORA_RANK,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        target_modules=LORA_TARGET_MODULES,
        bias="none",
        inference_mode=False,
    )
    model = get_peft_model(model, config)
    model.print_trainable_parameters()
    return model


def adapter_base_model(adapter_dir: str) -> Optional[str]:
    """Read base_model_name_or_path from the adapter's config, if available."""
    cfg_path = os.path.join(adapter_dir, "adapter_config.json")
    try:
        with open(cfg_path) as fh:
            return json.load(fh).get("base_model_name_or_path")
    except Exception:
        return None


def load_existing_lora(model, adapter_dir: str, cb: Optional[Callable] = None):
    from peft import PeftModel
    model = _prep_kbit(model)
    log(f"Loading existing LoRA adapter from {adapter_dir}", cb)
    model = PeftModel.from_pretrained(model, adapter_dir, is_trainable=True)
    model.print_trainable_parameters()
    return model

# ── Training ──────────────────────────────────────────────────────────────────

def train(cb: Optional[Callable] = None) -> str:
    from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling
    from datasets import Dataset as HFDataset

    log("=" * 55, cb)
    log("  MamunAI LoRA Fine-Tuning — Owner: Al Mamun Khan", cb)
    log("=" * 55, cb)

    # Auth
    hf_token = os.environ.get("HF_TOKEN", "").strip()
    if not hf_token:
        raise EnvironmentError(
            "HF_TOKEN is not set. Add it as a Space secret or environment variable."
        )
    try:
        login(token=hf_token)
    except Exception as exc:
        raise RuntimeError(f"Hugging Face authentication failed: {exc}") from exc
    create_repo(repo_id=HF_REPO, repo_type="model", exist_ok=True)

    # Dataset
    raw_data = download_dataset(cb)

    # Adapter check
    log(f"Checking Hugging Face repository: {HF_REPO}", cb)
    resume = adapter_exists_on_hub(HF_REPO, token=hf_token)

    if resume:
        log("Existing adapter found. Resuming training.", cb)
        if os.path.exists(OUTPUT_DIR):
            shutil.rmtree(OUTPUT_DIR)
        if not download_adapter_from_hub(HF_REPO, token=hf_token, local_dir=OUTPUT_DIR):
            log("Adapter download incomplete — falling back to fresh start.", cb)
            resume = False

    # Guard: if the downloaded adapter was trained on a different base model,
    # loading it would cause a size-mismatch crash. Detect this early and
    # start fresh with a clear message instead.
    if resume:
        saved_base = adapter_base_model(OUTPUT_DIR)
        if saved_base and saved_base != BASE_MODEL_ID:
            log(
                f"WARNING: Existing adapter was trained on '{saved_base}', "
                f"but current base model is '{BASE_MODEL_ID}'. "
                "Adapter is incompatible — discarding and starting fresh.",
                cb,
            )
            shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
            resume = False

    if not resume:
        log("No adapter found. Starting fresh.", cb)

    # Model
    model, tokenizer = load_base_model_and_tokenizer(cb)

    if resume and os.path.isfile(os.path.join(OUTPUT_DIR, "adapter_config.json")):
        model = load_existing_lora(model, OUTPUT_DIR, cb)
    else:
        model = apply_lora_fresh(model, cb)

    # Disable KV cache before gradient checkpointing
    model.config.use_cache = False
    if hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()
        log("Gradient checkpointing enabled.", cb)

    # Tokenise
    formatted = [format_example(r) for r in raw_data]

    def tokenize_fn(batch):
        out = tokenizer(
            batch["text"],
            truncation=True,
            max_length=MAX_SEQ_LENGTH,
            padding="max_length",
            return_tensors=None,
        )
        out["labels"] = out["input_ids"].copy()
        return out

    hf_ds    = HFDataset.from_dict({"text": formatted})
    tokenized = hf_ds.map(tokenize_fn, batched=True, remove_columns=["text"])
    log(f"Tokenized {len(tokenized)} examples.", cb)

    # Training args
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM,
        warmup_steps=WARMUP_STEPS,
        learning_rate=LEARNING_RATE,
        fp16=USE_FP16,
        bf16=USE_BF16,
        logging_steps=10,
        save_strategy="steps",
        save_steps=SAVE_STEPS,
        save_total_limit=1,
        optim="adamw_torch",
        report_to="none",
        dataloader_num_workers=0,
        remove_unused_columns=False,
        label_names=["labels"],
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=tokenized,
        data_collator=DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=False, pad_to_multiple_of=8
        ),
    )

    log("Starting training...", cb)
    trainer.train()

    # Save clean adapter (remove any checkpoint subdirs)
    log(f"Saving updated LoRA adapter to {OUTPUT_DIR}/", cb)
    for item in Path(OUTPUT_DIR).iterdir():
        if item.is_dir() and item.name.startswith("checkpoint-"):
            shutil.rmtree(item)
    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)

    # Cumulative metadata
    meta_path = os.path.join(OUTPUT_DIR, "training_meta.json")
    prev = {}
    if os.path.isfile(meta_path):
        with open(meta_path) as fh:
            prev = json.load(fh)
    runs  = prev.get("total_runs", 0) + 1
    total = prev.get("total_samples_trained", 0) + len(raw_data)
    with open(meta_path, "w") as fh:
        json.dump({
            "base_model":             BASE_MODEL_ID,
            "hf_repo":                HF_REPO,
            "lora_rank":              LORA_RANK,
            "lora_alpha":             LORA_ALPHA,
            "lora_dropout":           LORA_DROPOUT,
            "max_seq_length":         MAX_SEQ_LENGTH,
            "num_epochs_per_run":     NUM_EPOCHS,
            "total_runs":             runs,
            "total_samples_trained":  total,
            "last_run_samples":       len(raw_data),
            "target_modules":         LORA_TARGET_MODULES,
            "dataset_source":         GITHUB_DATASET_URL,
        }, fh, indent=2)
    log(f"Run #{runs} | samples this run: {len(raw_data)} | total: {total}", cb)

    # Upload
    log("Uploading updated adapter to Hugging Face.", cb)
    try:
        model.push_to_hub(HF_REPO, commit_message=f"Run {runs} — {len(raw_data)} samples")
        tokenizer.push_to_hub(HF_REPO, commit_message=f"Run {runs} — tokenizer")
        HfApi().upload_file(
            path_or_fileobj=meta_path,
            path_in_repo="training_meta.json",
            repo_id=HF_REPO, repo_type="model", token=hf_token,
            commit_message=f"Run {runs} — metadata",
        )
        log(f"Adapter uploaded → https://huggingface.co/{HF_REPO}", cb)
    except Exception as exc:
        log(f"Upload failed: {exc}  (adapter saved locally at {OUTPUT_DIR}/)", cb)

    log("Training completed successfully.", cb)
    return (
        f"Training completed successfully. "
        f"Run #{runs} | {len(raw_data)} examples | Total samples trained: {total}"
    )

# ── Merge ─────────────────────────────────────────────────────────────────────

def merge(
    adapter_dir: str = OUTPUT_DIR,
    merged_dir:  str = MERGED_DIR,
    upload_repo: Optional[str] = None,
    cb: Optional[Callable] = None,
) -> str:
    """
    Merge the LoRA adapter into the base model.
    Never called automatically — only when explicitly triggered.
    Optionally uploads the merged model to a separate HF repo.
    """
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from peft import PeftModel

    log("=" * 55, cb)
    log("  Merging LoRA adapter into base model...", cb)
    log("=" * 55, cb)

    cfg = os.path.join(adapter_dir, "adapter_config.json")

    if not os.path.isfile(cfg):
        log("Local adapter missing. Downloading from Hugging Face...", cb)

        hf_token = os.environ.get("HF_TOKEN", "").strip()

        download_adapter_from_hub(
            HF_REPO,
            token=hf_token,
            local_dir=adapter_dir
        )

    if not os.path.isfile(cfg):
        raise FileNotFoundError(
            f"Adapter not found locally or on HF repo: {HF_REPO}"
        )

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
        device_map="cpu"
    )

    log(f"Loading LoRA adapter from {adapter_dir}", cb)
    merged = PeftModel.from_pretrained(base, adapter_dir).merge_and_unload()

    os.makedirs(merged_dir, exist_ok=True)
    log(f"Saving merged model to {merged_dir}/", cb)
    merged.save_pretrained(merged_dir)
    tokenizer.save_pretrained(merged_dir)

    result = f"Merge completed successfully. Saved to {merged_dir}/"

    if upload_repo:
        hf_token = os.environ.get("HF_TOKEN", "").strip()
        if not hf_token:
            log("HF_TOKEN not set — skipping upload.", cb)
        else:
            try:
                login(token=hf_token)
                create_repo(repo_id=upload_repo, repo_type="model", exist_ok=True)
                log(f"Uploading merged model to {upload_repo}...", cb)
                merged.push_to_hub(upload_repo, commit_message="MamunAI merged model")
                tokenizer.push_to_hub(upload_repo, commit_message="MamunAI merged tokenizer")
                log(f"Merged model uploaded → https://huggingface.co/{upload_repo}", cb)
                result += f" | Uploaded to {upload_repo}"
            except Exception as exc:
                log(f"Upload failed: {exc}", cb)

    log("Merge completed successfully.", cb)
    return result

# ── CLI ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    p = argparse.ArgumentParser(description="MamunAI LoRA Fine-Tuning")
    p.add_argument("--merge",            action="store_true",
                   help="Merge saved adapter into base model instead of training")
    p.add_argument("--adapter-dir",      default=OUTPUT_DIR)
    p.add_argument("--merged-dir",       default=MERGED_DIR)
    p.add_argument("--upload-merged-to", default=None,
                   help="HF repo for the merged model upload (optional)")
    a = p.parse_args()

    if a.merge:
        merge(adapter_dir=a.adapter_dir, merged_dir=a.merged_dir, upload_repo=a.upload_merged_to)
    else:
        train()