"""Cloud GPU QLoRA training. Never falls back to training on the user's CPU."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil

from .data import encode_example, normalize_messages, prompt_key
from .runtime import CONFIG


def read_data(folder):
    manifest = json.loads((folder / "manifest.json").read_text(encoding="utf-8"))
    if manifest["config"] != CONFIG:
        raise ValueError("Data manifest does not match the pinned experiment configuration.")
    if "article_examples_sha256" in manifest:
        current = hashlib.sha256(Path(__file__).with_name("article_examples.json").read_bytes()).hexdigest()
        if manifest["article_examples_sha256"] != current:
            raise ValueError("Article behavior examples changed after preparation.")
    records, seen, groups = {}, set(), {}
    for split in ["train", "validation", "test"]:
        raw = (folder / f"{split}.jsonl").read_bytes()
        if hashlib.sha256(raw).hexdigest() != manifest["sha256"][split]:
            raise ValueError(f"Data changed after preparation: {split}")
        records[split] = [json.loads(line) for line in raw.decode().splitlines() if line]
        if not records[split] or len(records[split]) != manifest["counts"][split]:
            raise ValueError(f"Unexpected split size: {split}")
        for row in records[split]:
            key = prompt_key(normalize_messages(row))
            if key != row["prompt_hash"] or key in seen:
                raise ValueError("Duplicate or inconsistent prompts across dataset splits.")
            seen.add(key)
            group = row.get("group_id", key)
            if group in groups and groups[group] != split:
                raise ValueError("A source packet appears in multiple dataset splits.")
            groups[group] = split
    return records


def collate(examples, pad_token_id):
    import torch
    length = max(len(item["input_ids"]) for item in examples)
    return {key: torch.tensor([item[key] + [pad] * (length - len(item[key])) for item in examples])
            for key, pad in [("input_ids", pad_token_id), ("attention_mask", 0), ("labels", -100)]}


def train(args):
    import torch
    if not torch.cuda.is_available():
        raise RuntimeError("Cloud CUDA GPU required. No training or model download was started.")
    from datasets import Dataset
    from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from transformers import Trainer, TrainingArguments, set_seed

    output = Path(args.output)
    if output.exists() and any(output.iterdir()) and not args.resume_from_checkpoint:
        raise FileExistsError("Use an empty output directory or explicitly resume a checkpoint.")
    records = read_data(Path(args.data))
    set_seed(CONFIG["seed"])
    tokenizer = AutoTokenizer.from_pretrained(
        CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False)
    tokenizer.padding_side = "right"
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    base = AutoModelForCausalLM.from_pretrained(
        CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False,
        dtype=dtype, device_map={"": torch.cuda.current_device()}, attn_implementation="sdpa",
        quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=dtype))
    base.config.use_cache = False
    base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True,
                                           gradient_checkpointing_kwargs={"use_reentrant": False})
    model = get_peft_model(base, LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05,
        target_modules="all-linear", bias="none", task_type=TaskType.CAUSAL_LM))
    model.peft_config["default"].base_model_name_or_path = CONFIG["base_model"]
    model.peft_config["default"].revision = CONFIG["base_revision"]
    datasets = {split: Dataset.from_list([
        encode_example(tokenizer, row["messages"], CONFIG["max_sequence_length"]) for row in records[split]
    ]) for split in ["train", "validation"]}
    settings = TrainingArguments(
        output_dir=str(output), num_train_epochs=1, max_steps=args.max_steps,
        per_device_train_batch_size=1, per_device_eval_batch_size=1,
        gradient_accumulation_steps=16, learning_rate=5e-5, warmup_ratio=0.05,
        lr_scheduler_type="cosine", logging_steps=5, eval_strategy="steps", eval_steps=25,
        save_strategy="steps", save_steps=25, save_total_limit=2,
        load_best_model_at_end=True, metric_for_best_model="eval_loss",
        greater_is_better=False, bf16=dtype == torch.bfloat16, fp16=dtype == torch.float16,
        gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
        optim="adamw_torch", report_to="none", seed=CONFIG["seed"],
        dataloader_num_workers=0, remove_unused_columns=False)
    trainer = Trainer(model=model, args=settings, processing_class=tokenizer,
        train_dataset=datasets["train"], eval_dataset=datasets["validation"],
        data_collator=lambda batch: collate(batch, tokenizer.pad_token_id))
    baseline = trainer.evaluate()
    trained = trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
    final = trainer.evaluate()
    destination = output / "adapter"
    trainer.save_model(str(destination))
    tokenizer.save_pretrained(destination)
    shutil.copy2(Path(args.data) / "manifest.json", destination / "data_manifest.json")
    report = {"config": CONFIG, "baseline_validation": baseline, "final_validation": final,
              "training": trained.metrics, "counts": {k: len(v) for k, v in records.items()},
              "status": "candidate_requires_chat_code_and_article_evaluation", "max_steps": args.max_steps}
    (destination / "training_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    (destination / "README.md").write_text(
        "---\nlicense: gpl-3.0\nlibrary_name: peft\npipeline_tag: text-generation\n"
        f"base_model: {CONFIG['base_model']}\nbase_model_relation: adapter\n"
        "datasets:\n- HuggingFaceTB/smoltalk\ntags:\n- lora\n- code\n- experimental\n---\n\n"
        "# PocketVision Chat, Code & Writing candidate\n\nText-only QLoRA adapter. "
        "This candidate has validation-loss measurements only; no claim of better chat, coding or factual article quality. "
        "See training_report.json and data_manifest.json. "
        "Training data: Smol-Magpie-Ultra, Smol-rewrite, Smol-summarize (Apache-2.0), authored fictional article behavior examples (GPL-3.0-only), and the Self-OSS-Instruct subset "
        "(original: bigcode/self-oss-instruct-sc2-exec-filter-50k, ODC-BY). "
        "Original base weights remain Apache-2.0; project code and adapter use GPL-3.0-only.\n",
        encoding="utf-8")
    for source in ["LICENSE", "LICENSE-BASE-MODEL.txt"]:
        shutil.copy2(Path(__file__).resolve().parent.parent / source, destination / source)
    print(f"Candidate saved to {destination}. Evaluate answers before publishing a release.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train text-only Chat & Code on a cloud CUDA GPU")
    parser.add_argument("--data", default="chatcode-data")
    parser.add_argument("--output", required=True, help="Prefer mounted persistent cloud storage")
    parser.add_argument("--max-steps", type=int, default=-1, help="25 for a smoke run; -1 for one epoch")
    parser.add_argument("--resume-from-checkpoint")
    args = parser.parse_args()
    if args.max_steps != -1 and args.max_steps < 25:
        parser.error("Use at least 25 steps so validation and checkpoint saving occur.")
    train(args)