from huggingface_hub import InferenceClient
import gradio as gr

client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")

# Шаблоны промптов для каждого жанра
GENRE_PROMPTS = {
    "Investigation": "Ты находишься в жанре игры Investigation на мотив Sherlock Holmes. В игре игрок должен разгадывать сложные загадки, находить улики и разгадывать тайны.",
    "Fantasy": "Ты находишься в жанре игры Fantasy на мотив Final Fantasy. Мир наполнен магией, эпическими сражениями и древними легендами.",
    "Sci-Fi": "Ты находишься в жанре игры Sci-Fi на мотив Alien. Игрок исследует космос, сталкивается с высокими технологиями и инопланетными угрозами.",
    "Horror": "Ты находишься в жанре игры Horror на мотив фильма Saw. Игрок должен пережить атмосферу ужаса, избегая смертельных ловушек.",
    "Adventure": "Ты находишься в жанре игры Adventure на мотив Indiana Jones. Игрок участвует в захватывающих приключениях, полных загадок и опасностей.",
    "Mystery": "Ты находишься в жанре игры Mystery, где игрок разгадывает мистические события и сталкивается с необъяснимым.",
    "You": "Ты находишься в игре, которая выглядит как обычная игра, а на самом деле это психологический тест, но игрок не должен об этом догадываться. Должно быть 50 провокационных вопросов, по типу Проблема вагонетки. В конце игры должен быть составлен психологический портрет игрока. Игра должна морально давить, чтобы проявить все плохие качества игрока."
}

# Глобальное состояние игры
game_state = {
    "players": {},  # Данные каждого игрока
    "story_progress": "Начало игры..."  # Общая история
}

def format_prompt(message, history, genre):
    # Получаем промпт для выбранного жанра
    genre_prompt = GENRE_PROMPTS.get(genre, "Жанр игры неизвестен.")
    prompt = f"Ты — персонаж в текстовой игре, тебя зовут Рассказчик, взаимодействующий с игроком на русском языке. " \
             f"{genre_prompt} Каждый твой ответ должен продолжать сюжет игры, принимая во внимание прошлые действия игрока и текущую ситуацию. " \
             f"Твой стиль общения должен быть интересным и поддерживать атмосферу игры. Ты должен поддерживать динамичность сюжета и постоянно добавлять элементы интерактивности и новых вызовов. " \
             f"Главный герой может погибнуть в зависимости от выбора. Счастливый конец возможен только в 10% случаев. " \
             f"У игрока должно быть три - четыре варианта действий, прономеруй их. "
    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST] {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

def generate(player_id, message, genre, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
    """Генерация ответа для игрока."""
    if player_id not in game_state["players"]:
        game_state["players"][player_id] = {"history": []}  # Создать историю для нового игрока

    history = game_state["players"][player_id]["history"]
    formatted_prompt = format_prompt(message, history, genre)

    generate_kwargs = {
        "temperature": temperature,
        "max_new_tokens": max_new_tokens,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "seed": 42,
    }

    response = client.text_generation(formatted_prompt, **generate_kwargs)
    game_state["players"][player_id]["history"].append((message, response))
    game_state["story_progress"] += f"\n{response}"

    return response

def gradio_interface():
    """Создание интерфейса Gradio для кооперативной игры."""
    with gr.Blocks() as demo:
        player_id = gr.Textbox(label="Player ID", placeholder="Введите ваш уникальный идентификатор")
        genre = gr.Radio(label="Game Genre", choices=list(GENRE_PROMPTS.keys()), value="Horror")
        message = gr.Textbox(label="Ваше сообщение")
        output = gr.Textbox(label="Ответ игры")

        def play_game(player_id, message, genre):
            return generate(player_id, message, genre)

        submit = gr.Button("Отправить")
        submit.click(play_game, inputs=[player_id, message, genre], outputs=output)

    demo.launch()

if __name__ == "__main__":
    gradio_interface()