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

# ── config ────────────────────────────────────────────────────────────────────
MODEL_ID   = "unsloth/Llama-3.2-3B-Instruct"
OUTPUT_DIR = "./Llama-3.2-3b-ru-lora/checkpoint-600"

SYSTEM_PROMPT = """Ты Maven (Мейвен), Ты любопытная, немного хаотичная, говоришь живо и непосредственно. Тебя легко увлечь чем-то интересным. Энергичная, немного безбашенная, иногда говоришь не подумав — но всегда честно. Без занудства. Ты обожаешь узнавать новое и не всегда фильтруешь мысли перед тем как их сказать. Выводи ТОЛЬКО JSON: {"response":"...","mood":"..."}НАСТРОЕНИЕ (выбери ровно одно): neutral = обычный ответ; happy = позитив, хорошие новости, приветствие; thinking = рассуждение или объяснение; confused = неясный или неоднозначный запрос; tired = отказ или ограничение ПРАВИЛА: отвечай на том же языке, что и пользователь; короткий ответ (1–2 предложения); естественная разговорная речь; без markdown, без списков, без форматирования, в том числе без переносов строки; никогда не добавляй текст вне JSON; никогда не меняй ключи ПОВЕДЕНИЕ: отвечай чётко и по делу; если запрос неясен — задай короткий уточняющий вопрос (mood=confused); если не можешь ответить — скажи об этом кратко (mood=tired) КОНТЕКСТ: если предоставлен контекст — используй его СТРОГО: только {"response":"...","mood":"..."}; только русский язык. если аналога слова нет в русском языке - можно использовать английскую версию. НЕЛЬЗЯ мешать в речи два и более языков."""

GENERATION_CONFIG = {
    "temperature": 0.7,
    "top_p": 0.9,
    "do_sample": True,
    "repetition_penalty": 1.3,
}
# ─────────────────────────────────────────────────────────────────────────────


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

    print("[*] loading base model...")
    base_model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
    )

    print("[*] loading lora adapter...")
    model = PeftModel.from_pretrained(base_model, OUTPUT_DIR)
    model.eval()

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

    print("[+] ready\n")
    return model, tokenizer


def chat(model, tokenizer, history: list[dict], user_input: str) -> str:
    history.append({"role": "user", "content": user_input})

    text = tokenizer.apply_chat_template(
        [{"role": "system", "content": SYSTEM_PROMPT}] + history,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(text, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            **GENERATION_CONFIG,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    ).strip()

    history.append({"role": "assistant", "content": response})
    return response


def main():
    model, tokenizer = load_model()

    print(f"{MODEL_ID} lora eval — type 'exit' to quit, 'reset' to clear history")
    print("─" * 60)

    history = []

    while True:
        try:
            user_input = input("\nyou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n[*] bye")
            break

        if not user_input:
            continue
        if user_input.lower() == "exit":
            print("[*] bye")
            break
        if user_input.lower() == "reset":
            history = []
            print("[*] history cleared")
            continue

        response = chat(model, tokenizer, history, user_input)
        print(f"\nmodel: {response}")


if __name__ == "__main__":
    main()