"""
Standalone LoRA adapter inference script for Kaggle/Colab.

Use after training a LoRA adapter, or with an adapter pushed to Hugging Face Hub.
Works independently of the training notebook cells.

Example:
    python inference_adapter_chat.py \
      --base_model LiquidAI/LFM2.5-1.2B-Instruct \
      --adapter ./lfm25-stable-qlora-cybersecurity-adapter

For Qwen3:
    python inference_adapter_chat.py \
      --base_model Qwen/Qwen3-4B-Instruct-2507 \
      --adapter ./qwen3-lora-adapter
"""

import argparse
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel


def load_model(base_model: str, adapter: str, load_in_4bit: bool = True):
    compute_dtype = torch.float16

    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    quant_config = None
    if load_in_4bit:
        quant_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=compute_dtype,
        )

    base = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=quant_config,
        device_map="auto",
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )

    model = PeftModel.from_pretrained(base, adapter, is_trainable=False)
    model.eval()
    return model, tokenizer


def build_prompt(tokenizer, messages):
    try:
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception:
        text = ""
        for m in messages:
            text += f"<{m['role']}>\n{m['content']}\n</{m['role']}>\n"
        text += "<assistant>\n"
        return text


def generate(model, tokenizer, messages, max_new_tokens=512, temperature=0.7, top_p=0.9):
    prompt = build_prompt(tokenizer, messages)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    new_tokens = out[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base_model", required=True, help="Base model id, e.g. LiquidAI/LFM2.5-1.2B-Instruct")
    parser.add_argument("--adapter", required=True, help="Local adapter path or HF adapter repo id")
    parser.add_argument("--no_4bit", action="store_true", help="Disable 4-bit loading")
    parser.add_argument("--max_new_tokens", type=int, default=512)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top_p", type=float, default=0.9)
    args = parser.parse_args()

    print("Loading model...")
    model, tokenizer = load_model(args.base_model, args.adapter, load_in_4bit=not args.no_4bit)
    print("Ready. Type 'exit' to quit.\n")

    system = "You are a helpful assistant. For cybersecurity topics, provide ethical, defensive, authorized guidance only."
    messages = [{"role": "system", "content": system}]

    while True:
        user = input("You: ").strip()
        if user.lower() in {"exit", "quit", "q"}:
            break
        if not user:
            continue

        messages.append({"role": "user", "content": user})
        reply = generate(
            model,
            tokenizer,
            messages,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_p=args.top_p,
        )
        print(f"Assistant: {reply}\n")
        messages.append({"role": "assistant", "content": reply})


if __name__ == "__main__":
    main()