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

# Используем Phi-3-mini-instruct
INTERVIEWER_MODEL = "microsoft/Phi-3-mini-4k-instruct"

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=80,
    temperature=0.7,
    top_p=0.9,
)

# Память о пользователе
user_sessions = {}

# 16 категорий MBTI (можно адаптировать под твои .json)
MBTI_CATEGORIES = [
    "ENFJ", "ENFP", "ENTJ", "ENTP",
    "ESFJ", "ESFP", "ESTJ", "ESTP",
    "INFJ", "INFP", "INTJ", "INTP",
    "ISFJ", "ISFP", "ISTJ", "ISTP"
]

def clean_question(text: str) -> str:
    """Очищает вопрос от инструкций."""
    text = text.strip().split("\n")[0]
    text = text.strip('"').strip("'")
    for bad in ["user:", "assistant:", "system:", "instruction"]:
        if bad.lower() in text.lower():
            text = text.split(bad)[-1].strip()
    if not text.endswith("?"):
        text += "?"
    return text

async def generate_next_question(user_id: str, user_text: str = "") -> dict:
    """
    Генерирует следующий вопрос по категории.
    """
    session = user_sessions.get(user_id, {
        "history": [],
        "category_index": 0,
        "completed": False
    })

    # Проверяем, не закончились ли категории
    if session["category_index"] >= len(MBTI_CATEGORIES):
        session["completed"] = True
        user_sessions[user_id] = session
        return {
            "question": None,
            "category": None,
            "completed": True
        }

    current_category = MBTI_CATEGORIES[session["category_index"]]

    # Составляем контекст (историю последних 2 ответов)
    history = session["history"][-4:]
    formatted_history = "\n".join(
        [f"Q{i//2+1}: {history[i]}\nA{i//2+1}: {history[i+1]}" for i in range(0, len(history), 2)]
    )

    prompt = (
        f"You are an MBTI interviewer.\n"
        f"Current personality category: {current_category}\n\n"
        f"Previous dialogue:\n{formatted_history}\n\n"
        f"Now generate one new, open-ended question (starting with What, Why, How, or When) "
        f"to assess this category.\n"
        f"Do NOT repeat or rephrase any previous question.\n"
        f"Return only the question text."
    )

    # Генерация вопроса
    loop = asyncio.get_event_loop()
    raw = await loop.run_in_executor(None, lambda: llm_pipe(prompt)[0]["generated_text"])
    question = clean_question(raw)

    # Уникальный ID вопроса
    q_id = str(uuid.uuid4())[:8]

    # Обновляем состояние пользователя
    session["history"].extend([question, user_text])
    session["category_index"] += 1
    user_sessions[user_id] = session

    return {
        "id": q_id,
        "question": question,
        "category": current_category,
        "completed": False
    }