import os
import sys
import sqlite3
import pandas as pd
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    TrainerCallback
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

def adapted_projection_types(peft_model) -> set:
    """Base-module suffixes that actually received a LoRA adapter."""
    return {name.rsplit(".lora_A", 1)[0].split(".")[-1]
            for name, _ in peft_model.named_modules() if name.endswith("lora_A")}

EXPECTED_TARGETS = {"qkv_proj", "o_proj", "gate_up_proj", "down_proj"}
MAX_LEN = 2560  # training sequence cap; longer examples are dropped so their SQL labels aren't truncated


class VRAMLoggerCallback(TrainerCallback):
    """Callback to print VRAM usage during training to monitor resource consumption."""
    def on_step_end(self, args, state, control, **kwargs):
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated(0) / (1024 ** 3)
            reserved = torch.cuda.memory_reserved(0) / (1024 ** 3)
            print(f" - [Step {state.global_step}] VRAM Allocated: {allocated:.2f} GB | Reserved: {reserved:.2f} GB")

def train():
    print("=" * 60)
    print("Starting Phi-3 Text-to-SQL QLoRA Fine-Tuning")
    print("=" * 60)
    
    # 1. Device Configuration and Hardware Checks
    if not torch.cuda.is_available():
        print("WARNING: CUDA is NOT available. Running on CPU is extremely slow and NOT recommended.")
        device = "cpu"
        use_bf16 = False
        use_fp16 = False
    else:
        device = "cuda"
        gpu_name = torch.cuda.get_device_name(0)
        print(f"CUDA Device Detected: {gpu_name}")
        
        # Native bf16 only on Ampere+ (compute capability >= 8). Turing (T4, sm_75) reports
        # is_bf16_supported()=True but EMULATES bf16 slowly, so gate on the real capability.
        major, _ = torch.cuda.get_device_capability()
        use_bf16 = major >= 8
        use_fp16 = not use_bf16
        print(f"Precision: bf16={use_bf16} | fp16={use_fp16} (sm_{major}x)")

    # Quantization backend: 4-bit QLoRA only if bitsandbytes actually imports. It is broken on
    # some Kaggle/CUDA images (missing CUDA binary; 'No module named triton.ops'). Otherwise fall
    # back to plain LoRA in bf16/fp16 -- a 3.8B model fits a 16GB T4 without 4-bit.
    use_4bit = False
    if device == "cuda":
        try:
            import bitsandbytes  # noqa: F401
            use_4bit = True
            print("bitsandbytes available -> using 4-bit QLoRA")
        except Exception as e:
            print(f"[info] bitsandbytes unavailable ({type(e).__name__}: {e}); "
                  f"falling back to {'bf16' if use_bf16 else 'fp16'} LoRA (no 4-bit).")

    # 2. Paths
    dataset_dir = "data/spider"
    train_file = os.path.join(dataset_dir, "train.jsonl")
    val_file = os.path.join(dataset_dir, "dev.jsonl")
    output_dir = "models/phi3-text-to-sql"
    adapter_dir = "models/phi3-text-to-sql-adapter"
    
    if not os.path.exists(train_file) or not os.path.exists(val_file):
        print("ERROR: Dataset files not found. Run dataset.py first.")
        sys.exit(1)

    # 3. Model load config: 4-bit quant when available, else full-precision (bf16/fp16) LoRA.
    compute_dtype = torch.bfloat16 if use_bf16 else torch.float16
    model_id = "microsoft/Phi-3-mini-4k-instruct"
    load_kwargs = dict(trust_remote_code=False, attn_implementation="eager")
    if use_4bit:
        print("Configuring 4-bit Quantization (QLoRA)...")
        load_kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=compute_dtype,
        )
        load_kwargs["device_map"] = "auto"
        from accelerate import Accelerator
        load_kwargs["device_map"] = {"": Accelerator().local_process_index}
    else:
        load_kwargs["torch_dtype"] = compute_dtype
        load_kwargs["low_cpu_mem_usage"] = True   # no device_map: the Trainer places it on the GPU

    # 4. Load Model and Tokenizer
    print(f"Loading base model: {model_id} (4bit={use_4bit})...")
    model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
    
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=False)
    
    # Configure padding tokens for Phi-3 (which uses EOS for padding natively in standard chat SFT)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # 5. Prepare Model for Peft Training (gradient checkpointing + input grads)
    if use_4bit:
        print("Preparing model for k-bit training...")
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    else:
        print("Enabling gradient checkpointing for full-precision LoRA...")
        model.gradient_checkpointing_enable()
        model.enable_input_require_grads()

    # 6. Configure LoRA
    # Phi-3 fuses q/k/v -> qkv_proj and gate/up -> gate_up_proj; targeting these plus o_proj
    # and down_proj adapts all 7 logical projections (the v1 config silently hit only 2).
    print("Configuring LoRA parameters...")
    peft_config = LoraConfig(
        r=16, lora_alpha=32,
        target_modules=["qkv_proj", "o_proj", "gate_up_proj", "down_proj"],
        lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    )
    
    model = get_peft_model(model, peft_config)
    got = adapted_projection_types(model)
    assert got == EXPECTED_TARGETS, (
        f"LoRA adapted {got}, expected {EXPECTED_TARGETS}. Phi-3 fuses q/k/v->qkv_proj and "
        "gate/up->gate_up_proj; the v1 config silently matched only o_proj+down_proj. "
        "Refusing to train a partial adapter.")
    model.print_trainable_parameters()   # expect ~0.5-0.9% now (was 0.12%)

    # 7. Load Datasets
    print("Loading formatted training datasets...")
    dataset = load_dataset("json", data_files={
        "train": train_file,
        "validation": val_file
    })

    # Drop examples longer than MAX_LEN so their SQL labels are never truncated
    # (~0.95% of Spider train exceeds 2560 tokens; corrupt labels are worse than fewer rows).
    def _fits(ex):
        return len(tokenizer.apply_chat_template(ex["messages"], tokenize=True)) <= MAX_LEN
    before = len(dataset["train"])
    dataset["train"] = dataset["train"].filter(_fits)
    dataset["validation"] = dataset["validation"].filter(_fits)
    print(f"Length filter (<= {MAX_LEN} tok): kept {len(dataset['train'])}/{before} train "
          f"(dropped {before - len(dataset['train'])}); dev now {len(dataset['validation'])}")
    
    # 8. Training Arguments
    # Memory-optimized: batch_size=1, gradient_accumulation=8 (eff. batch size 8), paged_adamw optimizer
    print("Setting up memory-optimized SFT training arguments...")
    training_args = SFTConfig(
        output_dir=output_dir,
        num_train_epochs=2,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=1,
        gradient_accumulation_steps=4,
        group_by_length=True,
        learning_rate=2e-4,
        weight_decay=0.01,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        logging_steps=1,
        eval_strategy="steps",
        eval_steps=200,
        save_strategy="steps",
        save_steps=200,
        save_total_limit=2,
        optim="paged_adamw_8bit" if use_4bit else "adamw_torch",
        bf16=use_bf16,
        fp16=use_fp16,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        max_grad_norm=0.3,
        report_to="none",
        ddp_find_unused_parameters=False,
        remove_unused_columns=True, # Allow Trainer to drop the raw 'messages' column after tokenization
        max_seq_length=2560
    )

    # 9. Initialize SFTTrainer
    # SFTTrainer natively formats 'messages' using the model's chat template
    print("Initializing SFTTrainer...")
    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        peft_config=None,
        processing_class=tokenizer,
        args=training_args,
        callbacks=[VRAMLoggerCallback()] if device == "cuda" else []
    )

    # 10. Execute Fine-Tuning
    print("Launching Fine-Tuning loop...")
    trainer.train()
    
    # 11. Save the Fine-Tuned Adapter
    print(f"Fine-tuning complete! Saving adapter to {adapter_dir}...")
    trainer.model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)
    print("Adapter saved successfully. Ready for evaluation and deployment!")

if __name__ == "__main__":
    train()