import os
import random
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer, util
from openai import OpenAI



# -----------------------------
# 1. ЗАГРУЗКА СПИСКА СЛОВ И ВЫЧИСЛЕНИЕ ЭМБЕДДИНГОВ
# -----------------------------
WORDS_FILE = "words.txt"
with open(WORDS_FILE, encoding="utf-8") as f:
    words = [line.strip() for line in f if line.strip()]

# Use sentence-transformers/all-MiniLM-L6-v2 for local embeddings
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
embeddings = embedding_model.encode(words, convert_to_tensor=True)

# Pick a random target word
target_word = random.choice(words)
print(target_word)
idx_target = words.index(target_word)
target_emb = embeddings[idx_target : idx_target + 1]

# Compute cosine similarities between target and all words
cosine_sim = util.cos_sim(target_emb, embeddings)[0]
sim_values = cosine_sim.cpu().numpy()

# Build a lookup dict: word -> similarity
similarity_dict = {w: float(sim_values[i]) for i, w in enumerate(words)}

# Sort words by similarity descending and build rank lookup
sorted_indices = np.argsort(-sim_values)  # indices sorted by decreasing similarity
N = len(words)
rank_dict = {}
for rank_zero_based, idx in enumerate(sorted_indices):
    word = words[idx]
    rank_dict[word] = rank_zero_based + 1  # 1-based rank

# -----------------------------
# 2. INIT NEBIUS CLIENT FOR Qwen/Qwen3-14B
# -----------------------------
# Ensure NEBIUS_API_KEY is set in the environment
client = OpenAI(
    base_url="https://api.studio.nebius.com/v1/",
    api_key=os.environ.get("NEBIUS_API_KEY"),
)

# -----------------------------
# 3. CALLBACK: PROCESS GUESS
# -----------------------------
def guess_callback(guess: str, history: list, rank_history: list):
    guess = guess.strip().lower()
    if guess not in rank_dict:
        # Re-render history even if guess invalid
        history_html = ""
        if history:
            lines = []
            for i, (w, r) in enumerate(zip(history, rank_history)):
                p = (r - 1) / (N - 1)
                red = int(255 * p)
                green = int(255 * (1 - p))
                color_hex = f"#{red:02x}{green:02x}00"
                lines.append(
                    f"<span style='font-weight:bold; color:{color_hex};'>"
                    f"{i+1}. {w} — {r}</span>"
                )
            history_html = "<br>".join(lines)
        return (
            f"Word “{guess}” not found in dictionary.",
            history,
            rank_history,
            history_html
        )

    r = rank_dict[guess]
    history = history + [guess]
    rank_history = rank_history + [r]

    msg = f"Index in sorted list for “{guess}”: {r} of {N}"
    if guess == target_word:
        msg += "\n\n🎉 Congratulations! You guessed the word!"

    # Build colored HTML history
    lines = []
    for i, (w, rank_val) in enumerate(zip(history, rank_history)):
        p = (rank_val - 1) / (N - 1)
        red = int(255 * p)
        green = int(255 * (1 - p))
        color_hex = f"#{red:02x}{green:02x}00"
        lines.append(
            f"<span style='font-weight:bold; color:{color_hex};'>"
            f"{i+1}. {w} — {rank_val}</span>"
        )
    history_html = "<br>".join(lines)

    return msg, history, rank_history, history_html

# -----------------------------
# 4. CALLBACK: ASK LLM FOR HELP
# -----------------------------
def help_callback(history: list, rank_history: list):
    if not history:
        return "Make at least one guess before asking for hints."

    attempts_info = "\n".join(
        [f"{i+1}. Word: “{w}”, Index: {rank_history[i]} of {N}"
         for i, w in enumerate(history)]
    )
    prompt = (
        "The user has the following guess history (word and its index in the sorted list):\n\n"
        f"{attempts_info}\n\n"
        "Without revealing the secret word, suggest new guesses or hints "
        "that might guide the user toward the correct answer."
    )

    completion = client.chat.completions.create(
        model="Qwen/Qwen3-14B",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.6
    )
    print(prompt)
    print('-'*80)
    print(completion.choices[0].message.content)
    return completion.choices[0].message.content

# -----------------------------
# 5. BUILD GRADIO INTERFACE
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("""
    <h1 style="text-align:center; font-family:Helvetica;">🎯 Guess The Word 🌟</h1>
    <p style="font-size:16px; text-align:center;">
      A random word has been chosen. Enter your guess and see its <em>index</em> 
      in the sorted list of similarities.  
      Colors transition from green (best rank) to red (worst rank).  
      If stuck, ask the LLM for hints!
    </p>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            text_input = gr.Textbox(
                label="Enter your guess",
                placeholder="e.g. fire",
                interactive=True
            )
            button_guess = gr.Button("Submit Guess", variant="primary")
            output_guess = gr.Textbox(
                label="Result",
                interactive=False
            )

            button_help = gr.Button("Ask LLM for Hints", variant="secondary")
            output_help = gr.Textbox(
                label="Hints from LLM",
                interactive=False
            )
        with gr.Column(scale=1):
            gr.Markdown("<h2 style='font-family:Helvetica;'>Guess History</h2>")
            history_display = gr.HTML()

    state_history = gr.State(value=[])
    state_rank_history = gr.State(value=[])

    # On guess, return (message, new history list, new rank list, colored HTML history)
    button_guess.click(
        fn=guess_callback,
        inputs=[text_input, state_history, state_rank_history],
        outputs=[output_guess, state_history, state_rank_history, history_display]
    )

    button_help.click(
        fn=help_callback,
        inputs=[state_history, state_rank_history],
        outputs=[output_help]
    )

if __name__ == "__main__":
    demo.launch()