# model_loader.py
import os
import torch
import gc
import json
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel

# CONFIG - change these if your repo name differs
BASE_MODEL = "meta-llama/Llama-3.2-3B"   # or local path if you prefer
ADAPTER = "iatechaways/crm_model"            # your LoRA adapter on HF

# Get HF token from environment
HF_TOKEN = os.getenv("HF_TOKEN")

# NOTE: For ZeroGPU we LOAD on CPU and move to GPU only during inference.
def load_model_cpu():
    """
    Load base model + adapter into CPU memory (device_map='cpu' or low mem). 
    Return a dict that holds references to model and tokenizer.
    """
    # load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True, token=HF_TOKEN)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # load base model on CPU to minimize early GPU allocation
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        device_map={"": "cpu"},
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
                token=HF_TOKEN
    )

    model = PeftModel.from_pretrained(
        base,
        ADAPTER,
        device_map={"": "cpu"},
        torch_dtype=torch.float32,
        is_trainable=False,
        low_cpu_mem_usage=True,
                token=HF_TOKEN
    )

    # don't create pipeline here; we'll call generate with model directly
    return {
        "tokenizer": tokenizer,
        "model": model,
        "base_name": BASE_MODEL,
        "adapter": ADAPTER
    }

def move_model_to_gpu(info: dict):
    """Move model to GPU and set dtype to float16 for inference speed."""
    if not torch.cuda.is_available():
        return
    model = info["model"]
    # move to cuda
    model.to("cuda")
    # if tokenizer/pipeline required, we keep tokenizer on CPU
    torch.cuda.empty_cache()

def move_model_to_cpu(info: dict):
    """Move model back to CPU & free GPU memory."""
    if not torch.cuda.is_available():
        return
    model = info["model"]
    model.to("cpu")
    torch.cuda.empty_cache()
    gc.collect()

def generate_with_model(info: dict, email_text: str, mail_type: str, client_data: dict, max_new_tokens: int = 200):
    """Run text generation using the loaded model (assumes model on device)."""
    tokenizer = info["tokenizer"]
    model = info["model"]
    device = next(model.parameters()).device

    # Prepare client context
    client_context = (
        f"Client Info:\n"
        f"Name: {client_data.get('name')}\n"
        f"Policy No: {client_data.get('policy_number')}\n"
        f"Plan: {client_data.get('indiafirst_life_insurance_plan')}\n"
        f"Status: {client_data.get('status')}\n"
        f"Premium: ₹{client_data.get('premium_amt')}\n"
        f"Due Date: {client_data.get('premium_due_date')}\n\n"
    )

    prompt = f"""Email Category: {mail_type}

### Instruction:
You are an LIC customer-service assistant. Respond politely and concisely (2–5 sentences).
Do NOT invent facts that are not present in the client data. If data is missing, ask the customer to provide correct details.

{client_context}

### Input:
{email_text}

### Response:
"""

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.6,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.12,
            num_beams=1
        )
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return decoded