# ---------------------------------------------------------
# Install (if needed, run once in your environment):
# !pip install langgraph gradio matplotlib
# ---------------------------------------------------------

import random
from typing import Literal
from typing_extensions import TypedDict

import gradio as gr
import matplotlib.pyplot as plt
from langgraph.graph import StateGraph, START, END

# ==============================
# 1. LangGraph Game Definition
# ==============================

class GameState(TypedDict):
    bet: int           # Player's bet (2–12)
    roll: int          # Dice roll result (2–12)
    outcome: str       # "win" or "loss"
    winnings: int      # Total winnings across rounds


def node_1(state: GameState) -> GameState:
    """
    Entry node: assumes 'bet' and current 'winnings' are already in the state.
    We don't change them here; we just pass the state along.
    """
    # Nothing to change here, but you could log or validate.
    return state


def node_2(state: GameState) -> GameState:
    """
    Win node: player wins if roll <= bet.
    Add $10 to winnings and mark outcome.
    """
    state["outcome"] = "win"
    state["winnings"] += 10
    return state


def node_3(state: GameState) -> GameState:
    """
    Loss node: player loses otherwise.
    Subtract $5 from winnings and mark outcome.
    """
    state["outcome"] = "loss"
    state["winnings"] -= 5
    return state


def determine_outcome(state: GameState) -> Literal["node_2", "node_3"]:
    """
    Conditional edge function:
    - Roll two six-sided dice
    - If roll <= bet -> go to node_2 (win)
    - Else -> go to node_3 (loss)
    """
    roll = random.randint(1, 6) + random.randint(1, 6)
    state["roll"] = roll

    bet = int(state["bet"])
    return "node_2" if roll <= bet else "node_3"


# Build LangGraph
builder = StateGraph(GameState)

builder.add_node("node_1", node_1)
builder.add_node("node_2", node_2)
builder.add_node("node_3", node_3)

builder.add_edge(START, "node_1")
builder.add_conditional_edges("node_1", determine_outcome)
builder.add_edge("node_2", END)
builder.add_edge("node_3", END)

graph = builder.compile()

# ==============================
# 2. Gradio Interface
# ==============================

def play_round(
    bet: int,
    current_winnings: int,
    round_no: int,
    history: list
):
    """
    Gradio callback:
    - Takes current bet, winnings, round number, and history
    - Runs a single round through LangGraph
    - Returns updated UI + state
    """
    # Build initial LangGraph state from current UI state
    state: GameState = {
        "bet": int(bet),
        "roll": 0,
        "outcome": "",
        "winnings": int(current_winnings),
    }

    # Invoke the LangGraph game for a single round
    new_state = graph.invoke(state)

    new_winnings = new_state["winnings"]
    roll = new_state["roll"]
    outcome = new_state["outcome"]

    round_no = round_no + 1
    history = history + [(round_no, new_winnings)]

    # Create winnings-over-time plot
    rounds = [r for r, _ in history]
    winnings_vals = [w for _, w in history]

    fig, ax = plt.subplots()
    ax.plot(rounds, winnings_vals, marker="o")
    ax.set_xlabel("Round")
    ax.set_ylabel("Total Winnings ($)")
    ax.set_title("Winnings Over Time")

    if outcome == "win":
        msg = f"🎲 Rolled **{roll}** — You **WIN** this round! (+$10)"
    else:
        msg = f"🎲 Rolled **{roll}** — You **LOSE** this round. (–$5)"

    msg += f"\n\nCurrent total winnings: **${new_winnings}**"

    return (
        msg,            # result_output
        new_winnings,   # winnings_output
        fig,            # history_plot
        new_winnings,   # state_winnings
        round_no,       # state_round
        history         # state_history
    )


def stop_game(
    current_winnings: int,
    round_no: int,
    history: list
):
    """
    Gradio callback for the Stop/Reset button:
    - Shows final winnings
    - Resets the internal state so a new game can start
    """
    # Make a simple "stopped" plot
    fig, ax = plt.subplots()
    if history:
        rounds = [r for r, _ in history]
        winnings_vals = [w for _, w in history]
        ax.plot(rounds, winnings_vals, marker="o")
    ax.set_xlabel("Round")
    ax.set_ylabel("Total Winnings ($)")
    ax.set_title("Game Stopped")

    msg = (
        f"🧾 Game stopped.\n\n"
        f"Final winnings: **${current_winnings}**.\n\n"
        f"Press **Place Bet** to start a new game."
    )

    # Reset internal state
    return (
        msg,    # result_output
        0,      # winnings_output reset
        fig,    # history_plot
        0,      # state_winnings
        0,      # state_round
        []      # state_history
    )


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🎲 LangGraph Dice Betting Game (with Gradio)

        - Roll two six-sided dice.
        - You **win** if the roll is **less than or equal** to your bet.
        - **Win:** +$10, **Loss:** –$5  
        - You can play multiple rounds and watch your winnings evolve over time.
        """
    )

    with gr.Row():
        bet_input = gr.Slider(
            minimum=2,
            maximum=12,
            step=1,
            value=7,
            label="Your bet (2–12, you win if roll ≤ bet)"
        )
        play_btn = gr.Button("Place Bet", variant="primary")
        stop_btn = gr.Button("Stop / Reset Game", variant="secondary")

    with gr.Row():
        result_output = gr.Markdown(label="Round Result")
        winnings_output = gr.Number(
            label="Total Winnings ($)",
            value=0,
            interactive=False
        )

    history_plot = gr.Plot(label="Winnings Over Time")

    # Internal session state
    state_winnings = gr.State(0)
    state_round = gr.State(0)
    state_history = gr.State([])

    # Wire buttons
    play_btn.click(
        fn=play_round,
        inputs=[bet_input, state_winnings, state_round, state_history],
        outputs=[
            result_output,
            winnings_output,
            history_plot,
            state_winnings,
            state_round,
            state_history,
        ],
    )

    stop_btn.click(
        fn=stop_game,
        inputs=[state_winnings, state_round, state_history],
        outputs=[
            result_output,
            winnings_output,
            history_plot,
            state_winnings,
            state_round,
            state_history,
        ],
    )

# Uncomment this line to run locally:
# demo.launch()