ui/gradio_app.py

"""Gradio game UI for Myco — LLM-powered forest adventure."""

import ast
import json
import copy
import random
import threading
import logging
from typing import Any, cast
import gradio as gr

from game.engine import (
    collect_current, eat_current,
    follow_whisper, hf_companion_status, myco_reply,
    pick_current, study_current, check_auto_pickup,
    get_myco_narrative, RARITY_WEIGHTS, POISONOUS,
    get_myco_log, _LAST_REPLY,
    companion_reply, discover_mushroom, 
)

from game.catalog import load_mushrooms
_CATALOG = None

def _get_catalog():
    global _CATALOG
    if _CATALOG is None:
        _CATALOG = list(load_mushrooms())
    return _CATALOG
    
from game.state import welcome_history

from models.mushroom import Mushroom

from ui.renderers import (
    dex_markdown, forest_scene, game_intro_markdown,
    game_status_markdown, hero_markdown, home_forest_markdown,
    mushroom_card, progress_markdown, world_map_markdown, bubble
)

from game.progression import current_area, next_area, next_mystery, revealed_mysteries

from ui.interface import render_shroom_tab


logger = logging.getLogger(__name__)
EMPTY_COLLECTION = None
START_POSITION   = (1, 1)

# ---------------------------------------------------------------------------
# Autonomous roam state — pure Python, no LLM for decisions
# ---------------------------------------------------------------------------
_ROAM_LOCK = threading.Lock()
_ROAM = {
    "position":  [1, 1],
    "visited":   set(),
    "thought":   "The forest is waking up...",
    "ticks":     0,
}

def _roam_next_pos(x, y):
    candidates = []
    for dx, dy in [(0,-1),(0,1),(-1,0),(1,0)]:
        nx, ny = x+dx, y+dy
        if 0 <= nx <= 2 and 0 <= ny <= 2:
            w = 1 if (nx,ny) in _ROAM["visited"] else 4
            candidates.extend([(nx,ny)]*w)
    return list(random.choice(candidates)) if candidates else [x, y]

def _autonomous_tick(current, collection, active_mushrooms, position):
    """Pure Python: move Myco, check pickup, return updated state."""
    with _ROAM_LOCK:
        x, y = _ROAM["position"]
        safe_mushrooms = list(active_mushrooms or [])
        coll           = list(collection or [])

        # Check pickup at current position first
        if safe_mushrooms:
            current, coll, safe_mushrooms = check_auto_pickup(
                x, y, current, coll, safe_mushrooms
            )

        # Move to next cell
        nx, ny = _roam_next_pos(x, y)
        _ROAM["position"] = [nx, ny]
        _ROAM["visited"].add((nx, ny))
        _ROAM["ticks"] += 1
        pos = (nx, ny)

        # Check pickup at new position
        if safe_mushrooms:
            current, coll, safe_mushrooms = check_auto_pickup(
                nx, ny, current, coll, safe_mushrooms
            )

    return (
        list(pos),
        forest_scene(current, coll, pos, active_mushrooms=safe_mushrooms),
        game_status_markdown(current, coll, pos),
        mushroom_card_from_state(current),
        current,
        coll,
        safe_mushrooms,
        gr.update(),
        dex_markdown(coll),
        progress_markdown(coll),
    )

# ---------------------------------------------------------------------------
# JSON / history helpers
# ---------------------------------------------------------------------------
def force_parse_json_or_literal(s: str) -> Any:
    s = s.strip()
    if not (s.startswith("[") or s.startswith("{")):
        return s
    try:
        return json.loads(s)
    except Exception:
        pass
    try:
        return ast.literal_eval(s)
    except Exception:
        pass
    try:
        return json.loads(s.replace("None","null").replace("True","true").replace("False","false"))
    except Exception:
        pass
    return s

def extract_clean_text(val: Any) -> str:
    if val is None: return ""
    if isinstance(val, str):
        v = val.strip()
        if (v.startswith("[") and v.endswith("]")) or (v.startswith("{") and v.endswith("}")):
            parsed = force_parse_json_or_literal(v)
            if isinstance(parsed, str): return parsed
            return extract_clean_text(parsed)
        return val
    if isinstance(val, dict):
        for k in ("text","content"):
            if k in val: return extract_clean_text(val[k])
        return " ".join(extract_clean_text(v) for v in val.values() if v)
    if isinstance(val, list):
        return "\n".join(extract_clean_text(i) for i in val if i)
    return str(val)

def _normalize_history(history):
    if not history: return []
    if isinstance(history, str): history = force_parse_json_or_literal(history)
    if not isinstance(history, (list, tuple)): history = [history]
    raw = []
    for item in history:
        if item is None: continue
        if isinstance(item, str): item = force_parse_json_or_literal(item)
        if isinstance(item, dict):
            role = item.get("role","assistant")
            content = item.get("content", item.get("text",""))
            raw.append({"role": role, "content": content})
        elif isinstance(item, (list,tuple)) and len(item)==2:
            u, b = item
            if u: raw.append({"role":"user","content": extract_clean_text(u)})
            if b: raw.append({"role":"assistant","content": extract_clean_text(b)})
        else:
            raw.append({"role":"assistant","content": extract_clean_text(item)})
    reconciled = []
    for msg in raw:
        role = str(msg.get("role","assistant")).lower()
        if role not in ("user","assistant","system"): role = "assistant"
        content = str(extract_clean_text(msg.get("content",""))).strip()
        if not content: continue
        if reconciled and reconciled[-1]["role"] == role:
            reconciled[-1]["content"] += f"\n{content}"
        else:
            reconciled.append({"role": role, "content": content})
    return reconciled

def _denormalize_history(history):
    msgs = _normalize_history(history)
    out = []
    i = 0
    while i < len(msgs):
        role = msgs[i].get("role","assistant")
        content = msgs[i].get("content","")
        if role == "user":
            bot = ""
            if i+1 < len(msgs) and msgs[i+1].get("role") == "assistant":
                bot = msgs[i+1].get("content","")
                i += 2
            else:
                i += 1
            out.append((content, bot))
        else:
            out.append(("", content))
            i += 1
    return out

# ---------------------------------------------------------------------------
# Mushroom card helper
# ---------------------------------------------------------------------------
def mushroom_card_from_state(current):
    if current is None: return mushroom_card(None)
    if isinstance(current, str): current = force_parse_json_or_literal(current)
    if not isinstance(current, dict): return mushroom_card(None)
    state = copy.copy(current)
    state.setdefault("name","Unknown Mushroom")
    state.setdefault("rarity","Common")
    state.setdefault("habitat","Unknown")
    state.setdefault("lore","No lore available.")
    try:
        m = Mushroom.from_dict(state)
    except Exception:
        return mushroom_card(None)
    return mushroom_card(m, state.get("clue"), state.get("secret"), state)

# ---------------------------------------------------------------------------
# Search — spawns 4-6 mushrooms at random grid positions
# ---------------------------------------------------------------------------
def search_forest(position, collection):
    pos  = _safe_pos(position)
    coll = list(collection or [])
    catalog = _get_catalog()

    # Spawn 4-6 mushrooms at random grid positions
    count  = random.randint(4, 6)
    used   = set()
    batch  = []
    for _ in range(count):
        if len(used) >= 9: break
        while True:
            mx, my = random.randint(0,2), random.randint(0,2)
            if (mx,my) not in used: break
        used.add((mx,my))
        weights = [RARITY_WEIGHTS.get(getattr(m,"rarity","Common"),12) for m in catalog]
        m = random.choices(catalog, weights=weights, k=1)[0]
        batch.append({
            "name":    m.name,
            "rarity":  m.rarity,
            "habitat": getattr(m,"habitat","Forest"),
            "lore":    getattr(m,"lore",""),
            "edible":  getattr(m,"edible","Unknown"),
            "magic":   getattr(m,"magic","Unknown"),
            "danger":  getattr(m,"danger","Unknown"),
            "poison":  "Yes" if m.name in POISONOUS else "No",
            "mush_x":  mx,
            "mush_y":  my,
            "picked":  "No",
            "studied": "No",
            "clue":    "Move Myco onto this mushroom to collect it!",
            "score_delta": str({"Common":10,"Rare":35,"Legendary":100}.get(m.rarity,10)),
        })

    current = {
        "name":          "Mushroom Patch",
        "rarity":        "Multiple",
        "poison":        "No",
        "game_over":     "No",
        "picked":        "No",
        "studied":       "No",
        "encounter":     f"{count} mushrooms sprouted! Move Myco to collect them!",
        "event_title":   "Spore Influx",
        "event_emoji":   "🍄",
        "mystery_title": "The Wrong Memory",
        "mystery_next":  "Keep collecting to reveal the next chapter.",
        "score_total":   str(sum(int(b.get("score_delta","10")) for b in batch
                               if b.get("poison")=="No")),
        "health":        "3",
        "reward_text":   f"Move Myco to pick up {count} mushrooms!",
        "active_mushrooms": batch,
    }

    history = list(welcome_history()) + [
        {"role":"assistant","content":f"🍄 Whoa! {count} mushrooms just sprouted across the clearing! Move me around to collect them — but watch for the poisonous ones!"}
    ]

    # The "Current Discovery" card shows batch[0]. Make current_mushroom match
    # that same mushroom (merged with the patch's game-state fields like
    # mystery_title/event_title/score_total/active_mushrooms), so the action
    # buttons (Collect/Pick/Study/Whisper/Eat) operate on what's actually shown
    # instead of the abstract "Mushroom Patch" placeholder.
    spotlight = batch[0] if batch else None
    current_mushroom_state = {**current, **spotlight} if spotlight else current

    return (
        forest_scene(current_mushroom_state, coll, pos, active_mushrooms=batch),
        game_status_markdown(current_mushroom_state, coll, pos),
        mushroom_card_from_state(current_mushroom_state),
        current_mushroom_state,
        batch,                          # active_mushrooms state
        _normalize_history(history),
        dex_markdown(coll),
        progress_markdown(coll),
    )

# ---------------------------------------------------------------------------
# Movement — collision detection built in, no LLM call
# ---------------------------------------------------------------------------
def move_player(position, direction, current, collection, active_mushrooms):
    x, y = _safe_pos(position)
    if direction == "north":   y = max(0, y-1)
    elif direction == "south": y = min(2, y+1)
    elif direction == "west":  x = max(0, x-1)
    elif direction == "east":  x = min(2, x+1)
    pos = (x, y)

    safe_mushrooms = list(active_mushrooms or [])
    coll           = list(collection or [])

    # Collision check — pure Python, no LLM
    current, coll, safe_mushrooms = check_auto_pickup(
        x, y, current, coll, safe_mushrooms
    )

    # Update roam tracker
    with _ROAM_LOCK:
        _ROAM["position"] = [x, y]
        _ROAM["visited"].add((x, y))

    return (
        pos,
        forest_scene(current, coll, pos, active_mushrooms=safe_mushrooms),
        game_status_markdown(current, coll, pos),
        mushroom_card_from_state(current),
        current,
        coll,
        safe_mushrooms,
        gr.update(),              # leave chat untouched
        dex_markdown(coll),
        progress_markdown(coll),
    )

# ---------------------------------------------------------------------------
# Other player actions
# ---------------------------------------------------------------------------
def ask_companion(message, history, current, collection, position):
    legacy = _denormalize_history(history)
    result = myco_reply(message, legacy, current, collection, position)
    if isinstance(result, (tuple,list)) and len(result)==2:
        return result[0], _normalize_history(result[1])
    return "", _normalize_history(result)

def collect_discovery(current, collection, history, position):
    legacy = _denormalize_history(history)
    updated_coll, updated_hist = collect_current(current, collection, legacy)
    pos = _safe_pos(position)
    if isinstance(current,str): current = force_parse_json_or_literal(current)
    return (
        updated_coll,
        forest_scene(current, updated_coll, pos),
        game_status_markdown(current, updated_coll, pos),
        dex_markdown(updated_coll),
        world_map_markdown(updated_coll),
        _normalize_history(updated_hist),
        progress_markdown(updated_coll),
    )

def pick_discovery(current, collection, history, position):
    legacy = _denormalize_history(history)
    updated_coll, updated_current, updated_hist = pick_current(current, collection, legacy)
    if isinstance(updated_current,str): updated_current = force_parse_json_or_literal(updated_current)
    pos = _safe_pos(position)
    return (
        updated_coll,
        updated_current,
        forest_scene(updated_current, updated_coll, pos),
        game_status_markdown(updated_current, updated_coll, pos),
        mushroom_card_from_state(updated_current),
        dex_markdown(updated_coll),
        world_map_markdown(updated_coll),
        _normalize_history(updated_hist),
        progress_markdown(updated_coll),
    )

def follow_story_whisper(current, collection, history, position):
    legacy = _denormalize_history(history)
    updated_current, updated_hist = follow_whisper(current, collection, legacy)
    if isinstance(updated_current,str): updated_current = force_parse_json_or_literal(updated_current)
    pos = _safe_pos(position)
    return (
        updated_current,
        forest_scene(updated_current, collection, pos),
        game_status_markdown(updated_current, collection, pos),
        mushroom_card_from_state(updated_current),
        _normalize_history(updated_hist),
    )

def study_discovery(current, history, collection, position):
    legacy = _denormalize_history(history)
    # study_current() returns (observation_text, updated_history) — the
    # mushroom itself doesn't change when studied, only the chat does.
    _observation, updated_hist = study_current(current, legacy)
    if isinstance(current, str):
        current = force_parse_json_or_literal(current)
    pos = _safe_pos(position)
    return (
        current,
        forest_scene(current, collection, pos),
        game_status_markdown(current, collection, pos),
        mushroom_card_from_state(current),
        _normalize_history(updated_hist),
    )

def block_unsafe_eating(current, history):
    legacy = _denormalize_history(history)
    _, updated_hist = eat_current(current, None, legacy)
    return _normalize_history(updated_hist)

def open_play_tab():
    return gr.update(selected="play")

def _safe_pos(position):
    if position is None: return START_POSITION
    x, y = position
    return (max(0,min(2,int(x))), max(0,min(2,int(y))))


# ---------------------------------------------------------------------------
# Guide text
# ---------------------------------------------------------------------------
GUIDE_MARKDOWN = """
## How to play

1. Go to **🎮 Play** and click **🎮 Search Clearing** to spawn mushrooms across the grid.
2. Use **⬆️⬇️⬅️➡️** to move Myco — walk over mushrooms to collect them automatically!
3. **Watch out** — stepping on a poisonous mushroom triggers Game Over!
4. Myco also **roams autonomously** every 3 seconds collecting mushrooms for you.
5. Use the action buttons to Study, Collect, Pick, or Follow Whisper.
6. Ask Myco anything in the chat.

**Scoring:** Common = 10 · Rare = 35 · Legendary = 100 · Poison = −25 + Game Over

**Poisonous:** Ghost Gill · Pepper Pixie · Ruby Knuckle · Clockwork Chanterelle
"""

# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------
def build_app():
    with gr.Blocks(title="Myco — Forest Adventure") as demo:
        gr.HTML("""<style>
            .gradio-container { background: #1a1a2e !important; }
            #forest-scene { contain: layout style; content-visibility: auto; }
            .progress-level, .progress-level-inner, .generating { display: none !important; }
        </style>""")

        # ── Shared state ──
        current_mushroom = gr.State(None)
        collection       = gr.State([])
        player_position  = gr.State(START_POSITION)
        active_mushrooms = gr.State([])


        gr.HTML(hero_markdown())

        timer = gr.Timer(5)
        myco_bubble = gr.HTML()
        timer.tick(
            fn=lambda: bubble(get_myco_narrative()),
            outputs=myco_bubble
        )

        # Autonomous roam timer — every 5 seconds, pure Python movement
        roam_timer = gr.Timer(5)

        with gr.Tabs(selected="home") as tabs:

            with gr.Tab("🏠 Home", id="home"):
                gr.HTML(home_forest_markdown())
                start_play_button = gr.Button("🌲 Enter the Forest", variant="primary", size="lg")

            with gr.Tab("🎮 Play", id="play"):
                gr.HTML(game_intro_markdown())
                scene = gr.HTML(
                    forest_scene(None, EMPTY_COLLECTION, START_POSITION),
                    elem_id="forest-scene",
                )

                with gr.Row():
                    with gr.Column(scale=1):
                        search_button = gr.Button("🎮 Search Clearing", variant="primary", size="lg")
                    with gr.Column(scale=1):
                        gr.Markdown("**Move Myco**")
                        with gr.Row():
                            gr.HTML("<div></div>")
                            north_button = gr.Button("⬆️", size="sm")
                            gr.HTML("<div></div>")
                        with gr.Row():
                            west_button = gr.Button("⬅️", size="sm")
                            pick_button = gr.Button("🍄 Pick", variant="primary", size="sm")
                            east_button = gr.Button("➡️", size="sm")
                        with gr.Row():
                            gr.HTML("<div></div>")
                            south_button = gr.Button("⬇️", size="sm")
                            gr.HTML("<div></div>")

                status = gr.Markdown(game_status_markdown(None, EMPTY_COLLECTION, START_POSITION))

                with gr.Row(equal_height=True, elem_id="discovery-card"):
                    with gr.Column(scale=2):
                        gr.Markdown("### 🍄 Current Discovery")
                        discovery = gr.Markdown(mushroom_card(None))
                        with gr.Row():
                            study_button   = gr.Button("🔍 Study",          variant="secondary")
                            collect_button = gr.Button("🧺 Collect",        variant="secondary")
                        with gr.Row():
                            whisper_button = gr.Button("🌌 Follow Whisper", variant="secondary")
                            eat_button     = gr.Button("🍽️ Eat?",           variant="stop")

                    with gr.Column(scale=3):
                        gr.Markdown("### 💬 Chat with Myco")
                        chat = gr.Chatbot(
                            value=cast(Any, _normalize_history(welcome_history())),
                            label="Myco", height=400, show_label=False,
                        )
                        with gr.Row():
                            prompt     = gr.Textbox(
                                placeholder="Ask Myco about this place...",
                                show_label=False, scale=4, lines=1,
                            )
                            ask_button = gr.Button("Ask 🍄", variant="primary", scale=1)

            with gr.Tab("📖 MycoDex"):
                with gr.Row():
                    progress = gr.Markdown(progress_markdown(EMPTY_COLLECTION))
                    dex      = gr.Markdown(dex_markdown(EMPTY_COLLECTION))
                with gr.Group(elem_classes="myco-image-2"):
                    gr.Image(value="assets/images/Myco-2.png",
                             show_label=False, container=False, interactive=False)

            with gr.Tab("🗺️ Map"):
                world_map = gr.HTML(world_map_markdown(EMPTY_COLLECTION))

            with gr.Tab("🧭 Guide"):
                gr.Markdown(GUIDE_MARKDOWN)
                with gr.Group(elem_classes="myco-image"):
                    gr.Image(value="assets/images/Myco.png",
                             show_label=False, container=False, interactive=False)

            with gr.Tab("🍄 Myco"):
                render_shroom_tab()

        
        # Streams Myco's full log: narrative ([NARRATIVE]), in-game action
        # reactions ([ACTION]), and chat replies ([CHAT]).
        myco_live = gr.Textbox(
            label="🍄 Myco's live log (narrative / action / chat)",
            value="🍄 Myco is listening...",
            lines=10,
            max_lines=20,
            autoscroll=True,
            interactive=False,
            elem_id="myco-live-log",
        )

        # Poll the running Myco log (narrative/action/chat) once per second.
        # NOTE: previously this used demo.load(fn=_stream_myco, ...) with an
        # infinite `while True` generator. An infinite generator from
        # demo.load permanently occupies a queue worker/stream for that
        # session — other events (button clicks) can then execute their
        # Python function fine (visible in server logs) but never get their
        # result delivered back to the browser. A periodic Timer avoids this.
        log_timer = gr.Timer(1)
        log_timer.tick(fn=get_myco_log, outputs=myco_live)
                
        # ── Shared output lists ──
        _move_outputs = [
            player_position, scene, status, discovery,
            current_mushroom, collection, active_mushrooms,
            chat, dex, progress,
        ]

        # ── Event wiring ──
        start_play_button.click(open_play_tab, outputs=[tabs])

        # Roam timer — autonomous movement, pure Python
        roam_timer.tick(
            fn=_autonomous_tick,
            inputs=[current_mushroom, collection, active_mushrooms, player_position],
            outputs=_move_outputs,
        )

        search_button.click(
            search_forest,
            inputs=[player_position, collection],
            outputs=[scene, status, discovery, current_mushroom,
                     active_mushrooms, chat, dex, progress],
            show_progress="hidden",
        )

        ask_button.click(
            ask_companion,
            inputs=[prompt, chat, current_mushroom, collection, player_position],
            outputs=[prompt, chat], show_progress="hidden",
        )
        prompt.submit(
            ask_companion,
            inputs=[prompt, chat, current_mushroom, collection, player_position],
            outputs=[prompt, chat], show_progress="hidden",
        )

        # Arrow buttons — each hardcoded direction, correct inputs
        for direction, btn in [
            ("north", north_button), ("south", south_button),
            ("west",  west_button),  ("east",  east_button),
        ]:
            btn.click(
                lambda pos, col, cur, am, d=direction: move_player(pos, d, cur, col, am),
                inputs=[player_position, collection, current_mushroom, active_mushrooms],
                outputs=_move_outputs,
                show_progress="hidden",
            )

        study_button.click(
            study_discovery,
            inputs=[current_mushroom, chat, collection, player_position],
            outputs=[current_mushroom, scene, status, discovery, chat],
            show_progress="hidden",
        )
        collect_button.click(
            collect_discovery,
            inputs=[current_mushroom, collection, chat, player_position],
            outputs=[collection, scene, status, dex, world_map, chat, progress],
            show_progress="hidden",
        )
        pick_button.click(
            pick_discovery,
            inputs=[current_mushroom, collection, chat, player_position],
            outputs=[
                collection, current_mushroom, scene, status, 
                discovery, dex, world_map, chat, progress, 
            ],
            show_progress="hidden",
        )
        whisper_button.click(
            follow_story_whisper,
            inputs=[current_mushroom, collection, chat, player_position],
            outputs=[current_mushroom, scene, status, discovery, chat],
            show_progress="hidden",
        )
        eat_button.click(
            block_unsafe_eating,
            inputs=[current_mushroom, chat],
            outputs=[chat],
            show_progress="hidden",
        )

 # Auto-spawn mushrooms when app first loads
        demo.load(
                fn=search_forest,
                inputs=[player_position, collection],
                outputs=[scene, status, discovery, current_mushroom,
                         active_mushrooms, chat, dex, progress],
            )
        
        return demo