import json
import os
import random
import re
from typing import Any

import gradio as gr
from huggingface_hub import InferenceClient


MODEL_ID = os.getenv("HF_MODEL_ID", "meta-llama/Llama-3.1-8B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")
TOTAL_ROUNDS = 10
LETTERS = ["A", "B", "C", "D"]

SYSTEM_PROMPT = """You are the game master for "I See", a playful multiple-choice I Spy quiz.
Each round has one clue and four possible answers. Exactly one answer is correct.
Keep clues short, imaginative, and fair. Keep answers family-friendly and concrete."""

FALLBACK_ROUNDS = [
    {
        "clue": "I see something that's hollow inside",
        "answer": "A tree with an empty trunk",
        "options": [
            "A tree with an empty trunk",
            "A mountain with a doorbell",
            "A cloud with pockets",
            "A river that forgot to flow",
        ],
        "explanation": "A hollow tree trunk is empty inside.",
    },
    {
        "clue": "I see something that tells time without saying a word",
        "answer": "A clock on the wall",
        "options": [
            "A quiet drum",
            "A clock on the wall",
            "A sleeping notebook",
            "A candle in a cup",
        ],
        "explanation": "A clock shows time visually.",
    },
    {
        "clue": "I see something that carries rain above your head",
        "answer": "An umbrella",
        "options": [
            "A ladder",
            "A suitcase",
            "An umbrella",
            "A mirror",
        ],
        "explanation": "An umbrella blocks rain overhead.",
    },
    {
        "clue": "I see something that opens with teeth but never bites",
        "answer": "A zipper",
        "options": [
            "A zipper",
            "A spoon",
            "A window",
            "A shoelace",
        ],
        "explanation": "A zipper has teeth that join and separate.",
    },
    {
        "clue": "I see something that gets smaller every time it helps",
        "answer": "A pencil",
        "options": [
            "A chair",
            "A pencil",
            "A plate",
            "A blanket",
        ],
        "explanation": "A pencil gets shorter as it is sharpened and used.",
    },
]


def get_client() -> InferenceClient | None:
    if not HF_TOKEN:
        return None
    return InferenceClient(model=MODEL_ID, token=HF_TOKEN)


def call_llama(messages: list[dict[str, str]], max_tokens: int = 1800) -> str:
    client = get_client()
    if not client:
        raise RuntimeError("HF_TOKEN is not configured.")

    response = client.chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.9,
        top_p=0.92,
    )
    return response.choices[0].message.content.strip()


def extract_json(text: str) -> dict[str, Any]:
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        raise ValueError("No JSON object found in model response.")
    return json.loads(match.group(0))


def clean_text(value: Any) -> str:
    return re.sub(r"\s+", " ", str(value)).strip()


def fallback_round(round_number: int) -> dict[str, Any]:
    data = dict(FALLBACK_ROUNDS[(round_number - 1) % len(FALLBACK_ROUNDS)])
    options = list(data["options"])
    random.shuffle(options)
    data["options"] = options
    return data


def validate_round(data: dict[str, Any]) -> dict[str, Any]:
    clue = clean_text(data["clue"])
    answer = clean_text(data["answer"])
    options = [clean_text(option) for option in data["options"]]
    explanation = clean_text(data.get("explanation", ""))

    if len(options) != 4 or answer not in options:
        raise ValueError("Round must have four options and include the answer.")

    random.shuffle(options)
    return {
        "clue": clue,
        "answer": answer,
        "options": options,
        "explanation": explanation or f"The answer is {answer}.",
    }


def generate_game() -> list[dict[str, Any]]:
    try:
        model_text = call_llama(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": (
                        f"Create all {TOTAL_ROUNDS} rounds for one I See demo game. "
                        "Each round needs a short I Spy-style clue, four answer options, "
                        "one correct answer that appears exactly in options, and one short explanation. "
                        "Use varied everyday objects. Do not repeat correct answers. "
                        "Return only JSON with this exact shape: "
                        '{"rounds":[{"clue":"I see something that...","answer":"An umbrella",'
                        '"options":["An umbrella","A chair","A candle","A spoon"],'
                        '"explanation":"The umbrella keeps rain off you."}]}'
                    ),
                },
            ]
        )
        data = extract_json(model_text)
        rounds = [validate_round(round_data) for round_data in data["rounds"][:TOTAL_ROUNDS]]
        if len(rounds) != TOTAL_ROUNDS:
            raise ValueError("Game must include exactly 10 rounds.")
        return rounds
    except Exception:
        return [fallback_round(round_number) for round_number in range(1, TOTAL_ROUNDS + 1)]


def option_text(round_data: dict[str, Any], index: int) -> str:
    option = round_data["options"][index]
    return f"{LETTERS[index]}\n\n{option}"


def render_question(state: dict[str, Any] | None) -> tuple[str, str, str, str, str, str, str]:
    if not state:
        empty = "Start Game"
        return "I See", "ROUND 0 / 10", "SCORE: 0", empty, empty, empty, empty

    round_data = state["round"]
    return (
        f'## I Spy...\n\n### "{round_data["clue"]}"',
        f"ROUND {state['round_number']} / {TOTAL_ROUNDS}",
        f"SCORE: {state['score']}",
        option_text(round_data, 0),
        option_text(round_data, 1),
        option_text(round_data, 2),
        option_text(round_data, 3),
    )


def start_game() -> tuple[dict[str, Any], str, str, str, str, str, str, str, str]:
    rounds = generate_game()
    state = {
        "round_number": 1,
        "score": 0,
        "answered": False,
        "rounds": rounds,
        "round": rounds[0],
    }
    return state, *render_question(state), "", gr.update(visible=False)


def choose(index: int, state: dict[str, Any] | None) -> tuple[dict[str, Any] | None, str, str, Any]:
    if not state:
        return state, "Start the game first.", "SCORE: 0", gr.update(visible=False)

    if state["answered"]:
        return state, "You already answered this round. Hit Next Round.", f"SCORE: {state['score']}", gr.update(visible=True)

    selected = state["round"]["options"][index]
    answer = state["round"]["answer"]
    correct = selected == answer
    state["answered"] = True

    if correct:
        state["score"] += 1
        result = f"Correct. The answer is {answer}."
    else:
        result = f"Not quite. The answer is {answer}."

    explanation = state["round"].get("explanation", "")
    if explanation:
        result = f"{result}\n\n{explanation}"

    button_label = "Show Results" if state["round_number"] >= TOTAL_ROUNDS else "Next Round"
    return state, result, f"SCORE: {state['score']}", gr.update(value=button_label, visible=True)


def next_round(state: dict[str, Any] | None) -> tuple[dict[str, Any] | None, str, str, str, str, str, str, str, str, Any]:
    if not state:
        return None, "I See", "ROUND 0 / 10", "SCORE: 0", "Start Game", "Start Game", "Start Game", "Start Game", "", gr.update(visible=False)

    if state["round_number"] >= TOTAL_ROUNDS:
        score = state["score"]
        message = f"Game complete.\n\nFinal score: {score} / {TOTAL_ROUNDS}"
        return state, message, "FINISHED", f"SCORE: {score}", "A", "B", "C", "D", "Start a new game to play again.", gr.update(visible=False)

    round_number = state["round_number"] + 1
    state.update(
        {
            "round_number": round_number,
            "answered": False,
            "round": state["rounds"][round_number - 1],
        }
    )
    return state, *render_question(state), "", gr.update(visible=False)


APP_CSS = """
body,
.gradio-container {
    background: #050505 !important;
    color: #f7f7f7 !important;
}
.gradio-container {
    max-width: 980px !important;
    min-height: 100vh;
    padding: 22px 18px 34px !important;
    font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}
#stage {
    max-width: 700px;
    margin: 0 auto;
}
#title {
    margin: 0 0 28px;
}
#title h1 {
    text-align: center;
    margin-bottom: 0;
    font-size: 1.85rem;
    font-weight: 850;
    letter-spacing: 0;
}
#round-pill {
    text-align: center;
    width: fit-content;
    margin: 0 auto 18px;
    padding: 6px 14px;
    border: 1px solid #262626;
    border-radius: 999px;
    color: #f2f2f2;
    font-weight: 700;
    letter-spacing: 0.15em;
    text-transform: uppercase;
    font-size: 0.72rem;
}
#score-pill {
    text-align: center;
    color: #f4f4f4;
    font-weight: 750;
    letter-spacing: 0.16em;
    text-transform: uppercase;
    font-size: 0.8rem;
    margin: 18px 0 18px;
}
#question {
    text-align: center;
    white-space: pre-line;
    min-height: 136px;
    display: flex;
    align-items: center;
    justify-content: center;
    flex-direction: column;
}
#question h2 {
    font-size: clamp(1.8rem, 4vw, 2.7rem);
    font-style: italic;
    margin: 0.2rem 0 0.7rem;
    line-height: 1.05;
    letter-spacing: 0;
}
#question h3 {
    max-width: 760px;
    margin: 0 auto;
    color: #f7f7f7;
    font-size: clamp(1.05rem, 2.1vw, 1.35rem);
    font-style: italic;
    line-height: 1.45;
    font-weight: 800;
    letter-spacing: 0;
}
#result {
    min-height: 92px;
    text-align: center;
    white-space: pre-line;
    color: #d8d8d8;
    line-height: 1.45;
    padding-top: 8px;
}
#controls {
    max-width: 700px;
    margin: 0 auto 20px;
}
#answers {
    max-width: 700px;
    margin: 0 auto;
}
#answers .gr-row {
    gap: 18px;
}
#answers button {
    min-height: 158px;
    padding: 24px 26px !important;
    border: 1px solid #2a2a2a !important;
    border-radius: 22px !important;
    background: #050505 !important;
    color: #f7f7f7 !important;
    box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.015);
    justify-content: flex-start !important;
    align-items: flex-start !important;
    text-align: left !important;
    white-space: pre-line !important;
    font-weight: 800 !important;
    line-height: 1.35 !important;
    transition: border-color 140ms ease, transform 140ms ease, background 140ms ease;
}
#answers button:hover {
    border-color: #f5f5f5 !important;
    background: #0b0b0b !important;
    transform: translateY(-1px);
}
#answers button:active {
    transform: translateY(0);
}
#answers button::first-line {
    font-size: 0.78rem;
    letter-spacing: 0.14em;
}
#start,
#next {
    min-height: 52px;
}
#start button,
#next button {
    border-radius: 999px !important;
    font-weight: 800 !important;
    min-height: 52px;
    background: #ff6a13 !important;
    border: 0 !important;
    color: white !important;
}
#next button {
    background: #f7f7f7 !important;
    color: #080808 !important;
}
#footnote {
    max-width: 700px;
    margin: 28px auto 0;
    text-align: center;
    color: #8f8f8f;
    font-size: 0.82rem;
}
@media (max-width: 680px) {
    .gradio-container {
        padding: 18px 12px 28px !important;
    }
    #question {
        min-height: 126px;
    }
    #answers .gr-row {
        flex-direction: column;
        gap: 12px;
    }
    #answers button {
        min-height: 118px;
        padding: 20px !important;
    }
}
"""


with gr.Blocks(title="I See", css=APP_CSS) as demo:
    state = gr.State()

    with gr.Column(elem_id="stage"):
        gr.Markdown("# I See", elem_id="title")
        round_label = gr.Markdown("ROUND 0 / 10", elem_id="round-pill")
        question = gr.Markdown("## I Spy...\n\n### Start a game to generate your first question.", elem_id="question")
        score = gr.Markdown("SCORE: 0", elem_id="score-pill")

    with gr.Row(elem_id="controls"):
        start = gr.Button("Start Game", variant="primary", elem_id="start")
        next_button = gr.Button("Next Round", visible=False, elem_id="next")

    with gr.Column(elem_id="answers"):
        with gr.Row():
            choice_a = gr.Button("A")
            choice_b = gr.Button("B")
        with gr.Row():
            choice_c = gr.Button("C")
            choice_d = gr.Button("D")

    result = gr.Markdown("", elem_id="result")

    gr.Markdown(
        "Powered by `meta-llama/Llama-3.1-8B-Instruct` on Hugging Face when `HF_TOKEN` is configured.",
        elem_id="footnote",
    )

    outputs = [state, question, round_label, score, choice_a, choice_b, choice_c, choice_d, result, next_button]
    start.click(start_game, outputs=outputs)
    next_button.click(next_round, inputs=state, outputs=outputs)
    choice_a.click(lambda current: choose(0, current), inputs=state, outputs=[state, result, score, next_button])
    choice_b.click(lambda current: choose(1, current), inputs=state, outputs=[state, result, score, next_button])
    choice_c.click(lambda current: choose(2, current), inputs=state, outputs=[state, result, score, next_button])
    choice_d.click(lambda current: choose(3, current), inputs=state, outputs=[state, result, score, next_button])


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
    )