import os
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training
)
from datasets import load_dataset

# --- CLOSET & OFFLINE SETTINGS ---
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["WANDB_DISABLED"] = "true"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

def finetune(
    model_path: str,
    dataset_path: str,
    output_dir: str = "./output",
    batch_size: int = 1,
    gradient_accumulation_steps: int = 4,
    learning_rate: float = 2e-4,
    epochs: int = 1,
    lora_r: int = 8,
    lora_alpha: int = 16,
    lora_dropout: float = 0.05,
    hf_token: str = None
):
    if hf_token:
        from huggingface_hub import login
        login(token=hf_token)

    print(f"Loading tokenizer from {model_path}...")
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print("Loading model in 4-bit (QLoRA) for 8GB VRAM optimization...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )

    # Prepare for kbit training
    model.gradient_checkpointing_enable()
    model = prepare_model_for_kbit_training(model)

    # LoRA Configuration
    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], # Adjusted for common local architectures
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM"
    )

    model = get_peft_model(model, config)
    model.print_trainable_parameters()

    print(f"Loading dataset from {dataset_path}...")
    # Supporting local json/jsonl datasets OR Hugging Face datasets
    if os.path.exists(dataset_path):
        dataset = load_dataset("json", data_files=dataset_path, split="train")
    else:
        print(f"Dataset path '{dataset_path}' not found locally. Attempting to load from Hugging Face...")
        dataset = load_dataset(dataset_path, split="train", token=hf_token)

    def tokenize_function(examples):
        # 1. Handle Chat/Message format (e.g., ultrachat_200k)
        if "messages" in examples:
            # Format as alternating User/Assistant turns
            texts = []
            for msg_list in examples["messages"]:
                formatted_text = ""
                for msg in msg_list:
                    role = msg['role'].capitalize()
                    content = msg['content']
                    formatted_text += f"{role}: {content}\n"
                texts.append(formatted_text)
            return tokenizer(texts, truncation=True, padding="max_length", max_length=512)

        # 2. Handle Orca format (system, question, response)
        if "system_prompt" in examples and "question" in examples:
            texts = [
                f"### System:\n{s}\n\n### Question:\n{q}\n\n### Response:\n{r}"
                for s, q, r in zip(examples["system_prompt"], examples["question"], examples["response"])
            ]
            return tokenizer(texts, truncation=True, padding="max_length", max_length=512)

        # 3. Handle Alpaca format (instruction, input, output)
        if "instruction" in examples and "output" in examples:
            texts = []
            for i, inp, o in zip(examples["instruction"], examples.get("input", [""] * len(examples["instruction"])), examples["output"]):
                text = f"### Instruction:\n{i}"
                if inp:
                    text += f"\n\n### Input:\n{inp}"
                text += f"\n\n### Response:\n{o}"
                texts.append(text)
            return tokenizer(texts, truncation=True, padding="max_length", max_length=512)

        # 4. Fallback: Search for generic text fields
        text_field = "text"
        if "text" not in examples:
            for field in ["instruction", "content", "response", "output"]:
                if field in examples:
                    text_field = field
                    break
        
        return tokenizer(examples[text_field], truncation=True, padding="max_length", max_length=512)

    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)

    training_args = TrainingArguments(
        output_dir=output_dir,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        num_train_epochs=epochs,
        logging_steps=10,
        save_strategy="steps",
        save_steps=100,
        evaluation_strategy="no",
        fp16=False,
        bf16=True, # Recommended for RTX 40 series
        optim="paged_adamw_8bit", # Optimize VRAM
        report_to="none"
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )

    print("Starting training...")
    trainer.train()
    
    print(f"Saving final adapter to {output_dir}...")
    model.save_pretrained(output_dir)
    print("Finetuning complete.")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="8GB VRAM Optimized QLoRA Finetuner")
    parser.add_argument("--model", type=str, required=True, help="Local path to the model")
    parser.add_argument("--dataset", type=str, required=True, help="Local path or HF ID to .json/.jsonl dataset")
    parser.add_argument("--output", type=str, default="./finetuned_model", help="Directory to save output")
    parser.add_argument("--hf-token", type=str, help="Hugging Face API token")
    args = parser.parse_args()

    finetune(args.model, args.dataset, args.output, hf_token=args.hf_token)