# core/interviewer.py
import random
import itertools
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

INTERVIEWER_MODEL = "f3nsmart/TinyLlama-MBTI-Interviewer-LoRA"

tokenizer = AutoTokenizer.from_pretrained(INTERVIEWER_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    INTERVIEWER_MODEL, torch_dtype="auto", device_map="auto"
)

llm_pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=70,
    temperature=0.6,
    top_p=0.9,
)

CATEGORIES = [
    "Introversion", "Extroversion",
    "Sensing", "Intuition",
    "Thinking", "Feeling",
    "Judging", "Perceiving"
]

session_state = {}

def init_session(user_id: str):
    session_state[user_id] = {"asked": [], "answers": {}, "iteration": 1}

def select_next_category(user_id: str):
    s = session_state[user_id]
    remaining = [c for c in CATEGORIES if c not in s["asked"]]
    if not remaining:
        return None
    next_cat = random.choice(remaining)
    s["asked"].append(next_cat)
    return next_cat

def build_prompt(category: str):
    # ✅ Новый, более "демонстративный" промпт:
    return (
        f"You are a friendly MBTI interviewer.\n"
        f"Ask one short, open-ended question that explores {category.lower()}.\n"
        f"Examples: 'What makes you feel most energized in social situations?'\n"
        f"Output only the question, without quotes, without explanations."
    )

def clean_question(text: str) -> str:
    """Удаляет инструкции и оставляет только вопрос."""
    text = text.strip()

    # убираем строки с 'ask', 'instruction' и т.п.
    bad_phrases = ["ask", "instruction", "output only", "question about", "you are"]
    for phrase in bad_phrases:
        if phrase.lower() in text.lower():
            # берём только часть после последнего примера знака '?'
            if '?' in text:
                text = text.split('?')[-1]
            else:
                text = text.replace(phrase, '')
    text = text.strip().strip('"').strip("'")
    if not text.endswith("?"):
        text += "?"
    return text

def generate_question(user_id: str) -> str:
    """Генерация нового вопроса по категории."""
    if user_id not in session_state:
        init_session(user_id)

    category = select_next_category(user_id)
    if not category:
        return "✅ All 8 categories completed."

    prompt = build_prompt(category)
    raw = llm_pipe(prompt)[0]["generated_text"]
    question = clean_question(raw)
    return f"({category}) {question}"