"""
Hayat Elixir AI V5 - Meditron3-70B QLoRA Fine-Tuning Script
ERROR-PROOF version — addresses all 16 known HF Space errors
February 2026

CORE RULES APPLIED:
- ERR-03/15: Single A100-80GB, gradient_checkpointing, batch_size=1
- ERR-06: total_memory (not total_mem)
- ERR-09: trl==0.9.6 API (params in SFTTrainer, not SFTConfig)
- ERR-11/12: Single GPU only — NO DDP, NO device_map="auto" for multi-GPU
- ERR-16: Flask health server on port 7860 (may return HTML from HF proxy)
"""

import os
import json
import time
import torch
import threading
from flask import Flask, jsonify
from datetime import datetime
from huggingface_hub import HfApi, login

# ========== CONFIGURATION ==========
MODEL_ID = "OpenMeditron/Meditron3-70B"
DATASET_PATH = "/app/hayat_writer_v5_MASTER.jsonl"
OUTPUT_DIR = "/app/output"
ADAPTER_REPO = "mostafa922/hayat-meditron3-70b-clinical-v5"

# Training hyperparameters (PROVEN working on A100-80GB)
LORA_R = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05
NUM_EPOCHS = 3
BATCH_SIZE = 1               # ERR-15: batch_size=1 is the ONLY safe option for 70B on 80GB
GRADIENT_ACCUMULATION = 8    # Effective batch size = 8
LEARNING_RATE = 2e-4
MAX_SEQ_LENGTH = 1024        # ERR-15: 1024 is safe for A100-80GB
WARMUP_RATIO = 0.03

# Tokens
HF_TOKEN = os.environ.get("HF_TOKEN", "")

# ========== HEALTH SERVER (ERR-16: Flask on 7860) ==========
app = Flask(__name__)
training_status = {"stage": "initializing", "progress": 0, "message": "Starting up..."}

@app.route("/")
def health():
    return jsonify(training_status)

def start_health_server():
    app.run(host="0.0.0.0", port=7860, debug=False, use_reloader=False)

# Start health server in background thread
health_thread = threading.Thread(target=start_health_server, daemon=True)
health_thread.start()
print(f"[{datetime.now()}] Health server started on port 7860")

# ========== GPU CHECK (ERR-06: total_memory not total_mem) ==========
def check_gpu():
    if not torch.cuda.is_available():
        raise RuntimeError("No CUDA GPU available!")

    gpu_count = torch.cuda.device_count()
    print(f"\n{'='*60}")
    print(f"GPU REPORT")
    print(f"{'='*60}")

    for i in range(gpu_count):
        props = torch.cuda.get_device_properties(i)
        # ERR-06 FIX: Use total_memory, NOT total_mem
        vram_gb = props.total_memory / (1024**3)
        print(f"  GPU {i}: {props.name} | {vram_gb:.1f} GB VRAM")

    # ERR-11/12: Warn if multiple GPUs (we must use single GPU)
    if gpu_count > 1:
        print(f"\n  ⚠️  {gpu_count} GPUs detected — using GPU 0 ONLY (QLoRA + BitsAndBytes = single GPU)")

    primary_vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    if primary_vram < 70:
        raise RuntimeError(f"GPU 0 has only {primary_vram:.1f}GB VRAM. Need ≥80GB for 70B QLoRA. Use a100-large flavor.")

    print(f"\n  ✅ GPU 0 has {primary_vram:.1f}GB — sufficient for 70B QLoRA")
    print(f"{'='*60}\n")
    return primary_vram

# ========== MAIN TRAINING ==========
def main():
    global training_status
    start_time = time.time()

    # Step 0: Login to HuggingFace
    training_status = {"stage": "authenticating", "progress": 5, "message": "Logging into HuggingFace..."}
    print(f"[{datetime.now()}] Authenticating with HuggingFace...")

    if HF_TOKEN:
        login(token=HF_TOKEN)
        print(f"  ✅ Authenticated with HF token")
    else:
        print(f"  ⚠️  No HF_TOKEN set — gated model access may fail")

    # Step 1: GPU Check
    training_status = {"stage": "gpu_check", "progress": 10, "message": "Checking GPU..."}
    vram_gb = check_gpu()

    # Step 2: Load model with QLoRA 4-bit
    training_status = {"stage": "loading_model", "progress": 15, "message": "Loading Meditron3-70B (4-bit quantized)..."}
    print(f"[{datetime.now()}] Loading {MODEL_ID} with 4-bit quantization...")

    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    # ERR-11/12 FIX: Single GPU only — device_map targets GPU 0 explicitly
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map={"": 0},  # Force everything to GPU 0
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        token=HF_TOKEN,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        token=HF_TOKEN,
    )

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    # Log VRAM after model load
    allocated = torch.cuda.memory_allocated(0) / (1024**3)
    reserved = torch.cuda.memory_reserved(0) / (1024**3)
    print(f"  VRAM after model load: {allocated:.1f}GB allocated / {reserved:.1f}GB reserved / {vram_gb:.1f}GB total")

    training_status = {"stage": "model_loaded", "progress": 35, "message": f"Model loaded. VRAM: {allocated:.1f}/{vram_gb:.1f}GB"}

    # Step 3: Prepare for QLoRA
    print(f"[{datetime.now()}] Preparing QLoRA adapter (r={LORA_R}, alpha={LORA_ALPHA})...")
    model = prepare_model_for_kbit_training(model)

    # ERR-15 FIX: gradient_checkpointing saves ~40% VRAM
    model.gradient_checkpointing_enable()

    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
    )

    model = get_peft_model(model, lora_config)
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"  Trainable parameters: {trainable_params:,} / {total_params:,} ({100*trainable_params/total_params:.2f}%)")

    # Step 4: Load dataset
    training_status = {"stage": "loading_data", "progress": 45, "message": "Loading 840-example dataset..."}
    print(f"[{datetime.now()}] Loading dataset from {DATASET_PATH}...")

    from datasets import Dataset

    examples = []
    with open(DATASET_PATH) as f:
        for line in f:
            d = json.loads(line)
            # Format as chat template
            messages = d["messages"]
            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
            examples.append({"text": text})

    dataset = Dataset.from_list(examples)
    print(f"  ✅ Loaded {len(dataset)} examples")

    # Step 5: Training
    training_status = {"stage": "training", "progress": 50, "message": "Training started (3 epochs)..."}
    print(f"\n[{datetime.now()}] Starting training...")
    print(f"  Config: epochs={NUM_EPOCHS}, batch={BATCH_SIZE}, grad_accum={GRADIENT_ACCUMULATION}")
    print(f"  Effective batch size: {BATCH_SIZE * GRADIENT_ACCUMULATION}")
    print(f"  Learning rate: {LEARNING_RATE}")
    print(f"  Max seq length: {MAX_SEQ_LENGTH}")

    from transformers import TrainingArguments
    # ERR-09 FIX: trl==0.9.6 — params go in SFTTrainer(), NOT SFTConfig()
    from trl import SFTTrainer

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION,
        learning_rate=LEARNING_RATE,
        warmup_ratio=WARMUP_RATIO,
        # ERR-15 FIX: 8-bit optimizer reduces optimizer state memory by 50%
        optim="paged_adamw_8bit",
        fp16=False,
        bf16=True,
        logging_steps=5,
        save_strategy="epoch",
        save_total_limit=2,
        # ERR-15 FIX: gradient checkpointing saves ~40% VRAM
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        report_to="none",
        max_grad_norm=0.3,
        lr_scheduler_type="cosine",
        seed=42,
    )

    # ERR-09 FIX: trl==0.9.6 API — max_seq_length & dataset_text_field go here
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
        max_seq_length=MAX_SEQ_LENGTH,
        dataset_text_field="text",
        packing=False,
    )

    # Custom callback to update status
    class StatusCallback:
        def on_log(self, args, state, control, logs=None, **kwargs):
            global training_status
            if state.global_step > 0 and logs:
                progress = min(95, 50 + int(45 * state.global_step / state.max_steps))
                loss = logs.get("loss", "N/A")
                training_status = {
                    "stage": "training",
                    "progress": progress,
                    "message": f"Step {state.global_step}/{state.max_steps} | Loss: {loss}",
                    "step": state.global_step,
                    "max_steps": state.max_steps,
                    "loss": loss,
                }
                print(f"  Step {state.global_step}/{state.max_steps} | Loss: {loss}")

    trainer.add_callback(StatusCallback())

    # Train!
    train_result = trainer.train()

    train_time = time.time() - start_time
    print(f"\n[{datetime.now()}] Training complete!")
    print(f"  Total time: {train_time/60:.1f} minutes")
    print(f"  Final loss: {train_result.training_loss:.4f}")

    # Step 6: Save adapter
    training_status = {"stage": "saving", "progress": 96, "message": "Saving adapter files..."}
    print(f"[{datetime.now()}] Saving adapter to {OUTPUT_DIR}...")

    trainer.save_model(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)

    # Save training summary
    summary = {
        "model_id": MODEL_ID,
        "adapter_repo": ADAPTER_REPO,
        "dataset_size": len(dataset),
        "training_time_minutes": round(train_time / 60, 1),
        "final_loss": round(train_result.training_loss, 4),
        "epochs": NUM_EPOCHS,
        "lora_r": LORA_R,
        "lora_alpha": LORA_ALPHA,
        "learning_rate": LEARNING_RATE,
        "max_seq_length": MAX_SEQ_LENGTH,
        "batch_size": BATCH_SIZE,
        "gradient_accumulation": GRADIENT_ACCUMULATION,
        "effective_batch_size": BATCH_SIZE * GRADIENT_ACCUMULATION,
        "optimizer": "paged_adamw_8bit",
        "gpu": torch.cuda.get_device_name(0),
        "vram_gb": round(vram_gb, 1),
        "version": "V5-840ex-compliance-fix",
        "timestamp": datetime.now().isoformat(),
        "compliance_fix": "Added 20 compliance trap examples (Q24/Q37 Arabic greeting fix)",
    }

    with open(os.path.join(OUTPUT_DIR, "training_summary.json"), "w") as f:
        json.dump(summary, f, indent=2)

    print(f"  ✅ Adapter files saved")

    # Step 7: Upload to HuggingFace Hub
    training_status = {"stage": "uploading", "progress": 97, "message": "Uploading adapter to HuggingFace Hub..."}
    print(f"[{datetime.now()}] Uploading adapter to {ADAPTER_REPO}...")

    api = HfApi()
    try:
        api.create_repo(repo_id=ADAPTER_REPO, exist_ok=True, token=HF_TOKEN)
        api.upload_folder(
            folder_path=OUTPUT_DIR,
            repo_id=ADAPTER_REPO,
            token=HF_TOKEN,
            commit_message=f"V5 retrain: 840 examples with compliance fix (Q24/Q37)",
        )
        print(f"  ✅ Adapter uploaded to https://huggingface.co/{ADAPTER_REPO}")
    except Exception as e:
        print(f"  ⚠️  Upload error: {e}")
        print(f"  Adapter files saved locally at {OUTPUT_DIR}")

    # Step 8: List output files
    print(f"\n{'='*60}")
    print(f"OUTPUT FILES")
    print(f"{'='*60}")
    for f in sorted(os.listdir(OUTPUT_DIR)):
        fpath = os.path.join(OUTPUT_DIR, f)
        size_mb = os.path.getsize(fpath) / (1024*1024)
        print(f"  {f}: {size_mb:.1f} MB")

    # Final status
    training_status = {
        "stage": "completed",
        "progress": 100,
        "message": f"Training complete! Loss: {train_result.training_loss:.4f} | Time: {train_time/60:.1f}min",
        "summary": summary,
    }

    print(f"\n{'='*60}")
    print(f"✅ TRAINING COMPLETE")
    print(f"  Dataset: {len(dataset)} examples (840 = 820 original + 20 compliance fixes)")
    print(f"  Final loss: {train_result.training_loss:.4f}")
    print(f"  Time: {train_time/60:.1f} minutes")
    print(f"  Adapter: {ADAPTER_REPO}")
    print(f"{'='*60}")

    # Keep alive for log reading (10 min then auto-exit)
    print(f"\n[{datetime.now()}] Keeping alive for 10 minutes for log reading...")
    print(f"  ⚠️  REMEMBER: Pause this Space immediately after downloading adapter!")
    time.sleep(600)
    print(f"[{datetime.now()}] Auto-exit. PAUSE THIS SPACE NOW to avoid billing drain!")

if __name__ == "__main__":
    main()