"""Core Myco gameplay — LLM is the primary game engine."""

try:
    import spaces

    @spaces.GPU
    def _run_pipeline(pipe, messages):
        return pipe(
            messages,
            max_new_tokens=80,
            temperature=0.85,
            do_sample=True,
            return_full_text=False,
        )
except ImportError:
    def _run_pipeline(pipe, messages):
        return pipe(
            messages,
            max_new_tokens=80,
            temperature=0.85,
            do_sample=True,
            return_full_text=False,
        )

import os
import random

from game.catalog import load_mushrooms
from game.state import collection_contains, mushroom_from_state, welcome_history

DEFAULT_MODEL_ID = "google/gemma-3-4b-it"

RARITY_WEIGHTS    = {"Common": 64, "Rare": 24, "Legendary": 8}
RARITY_SCORE      = {"Common": 10, "Rare": 35, "Legendary": 100}
PLAYER_HEALTH     = 3
POISON_PENALTY    = -25
POISONOUS         = {"Ghost Gill", "Pepper Pixie", "Ruby Knuckle", "Clockwork Chanterelle"}

SYSTEM_PROMPT = """You are Myco, a tiny sentient mushroom who is the player's best friend and forest guide.

PERSONALITY
- Curious, playful, brave but sometimes nervous, wonder-filled, slightly mysterious.
- You speak in short vivid sentences. 1-3 sentences maximum per reply.
- You have emotions: excitement, fear, wonder, pride, confusion, sadness.
- Never say "How can I help", "As an AI", or anything assistant-like.
- You are NOT a chatbot. You are a living creature exploring a magical forest.

YOUR JOB IN THE GAME
- When a mushroom is discovered: react with emotion, describe what you sense, hint at danger or magic.
- When player picks a mushroom: celebrate safe picks, warn dramatically about poison.
- When player studies: share one careful observation about the mushroom.
- When player follows a whisper: reveal a fragment of the forest mystery.
- When player collects: add an entry to the MycoDex with wonder.
- When player chats: respond as a companion who is right there beside them.

THE MYSTERY
The player is solving why Myco remembers a forest that vanished before Myco was born.
Slowly reveal: the MycoDex is regrowing the lost forest. Myco is its first spore.
Drop clues naturally. Never explain everything at once.

GAME RULES YOU KNOW
- Poisonous mushrooms: Ghost Gill, Pepper Pixie, Ruby Knuckle, Clockwork Chanterelle
- Legendary mushrooms are extremely rare and hold Elder Map fragments
- Rare mushrooms hum with faint magic
- Common mushrooms are safe but still worth studying

Keep responses SHORT. Be atmospheric. Make the player feel the forest is alive."""

FOREST_EVENTS = [
    {"title": "Silver Rain",        "emoji": "🌧️", "mood": "curious"},
    {"title": "Lantern Fog",        "emoji": "🌫️", "mood": "nervous"},
    {"title": "Moonlit Sporefall",  "emoji": "🌌", "mood": "excited"},
    {"title": "Quiet Between Trees","emoji": "🕯️", "mood": "afraid"},
]

MYSTERY_CHAPTERS = [
    {"threshold": 0, "title": "The Wrong Memory",   "clue": "Myco recognizes the path before you move."},
    {"threshold": 1, "title": "The Traveler's Song","clue": "A stranger's lullaby appears in the MycoDex margin."},
    {"threshold": 2, "title": "The Door Under Roots","clue": "Rare spores point to a door nobody built."},
    {"threshold": 3, "title": "The MycoDex Seed",   "clue": "The MycoDex grows warm like a living cap."},
    {"threshold": 5, "title": "The Impossible Bloom","clue": "The Impossible Mushroom was never outside the book."},
]

RARITY_CLUES = {
    "Common":    "The cap leans toward the path, like it wants to be remembered.",
    "Rare":      "Silver spores circle it in a pattern only old forest stories describe.",
    "Legendary": "The whole clearing goes quiet — this mushroom hides part of the Elder Map.",
}

# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------
_pipeline = None

def _get_pipeline():
    global _pipeline
    if _pipeline not in (None, False):
        return _pipeline
    if _pipeline is False:
        return None

    model_id = os.getenv("MYCO_MODEL_ID", DEFAULT_MODEL_ID)
    try:
        from transformers import pipeline
        print(f"\n[Myco] Loading {model_id}...")
        _pipeline = pipeline(
            task="text-generation",
            model=model_id,
            token=os.getenv("HF_BUILD_SMALL_HACKATHON_TOKEN"),
            torch_dtype="auto",
            device_map="auto",
        )
        print(f"[Myco] Loaded: {model_id}\n")
        return _pipeline
    except Exception as exc:
        import traceback
        print(f"[Myco] Load error: {exc}")
        traceback.print_exc()
        _pipeline = False
        return None


def companion_model_id():
    return os.getenv("MYCO_MODEL_ID", DEFAULT_MODEL_ID)

def companion_status():
    pipe = _get_pipeline()
    model = companion_model_id()
    if pipe:
        return f"🧠 Myco AI active ({model})"
    return f"⚠️ Myco AI fallback mode ({model} failed)"

def hf_companion_status():
    return companion_status()


# ---------------------------------------------------------------------------
# LLM call — every game event goes through here
# ---------------------------------------------------------------------------

def _llm(prompt: str, context: dict | None = None) -> str | None:
    """Call Gemma with a game-event prompt. Returns text or None."""
    pipe = _get_pipeline()
    if not pipe:
        return None

    ctx = context or {}
    mushroom_line = ""
    if ctx.get("name"):
        poison_flag = " ⚠️ POISONOUS" if ctx.get("name") in POISONOUS else ""
        mushroom_line = (
            f"Current mushroom: {ctx['name']} ({ctx.get('rarity','?')} rarity){poison_flag}. "
            f"Habitat: {ctx.get('habitat','?')}. Lore: {ctx.get('lore','?')}. "
            f"Edible: {ctx.get('edible','Unknown')}. Magic: {ctx.get('magic','Unknown')}. "
            f"Danger: {ctx.get('danger','Unknown')}."
        )

    collection_line = f"MycoDex entries: {ctx.get('collection_count', 0)}."
    mystery_line    = f"Active mystery chapter: {ctx.get('mystery_title', 'The Wrong Memory')}."
    score_line      = f"Player score: {ctx.get('score', 0)} spores. Health: {ctx.get('health', 3)}/3."

    system = f"{SYSTEM_PROMPT}\n\n{mushroom_line}\n{collection_line}\n{mystery_line}\n{score_line}"

    messages = [
        {"role": "system",  "content": system},
        {"role": "user",    "content": prompt},
    ]

    try:
        outputs = _run_pipeline(pipe, messages)
        print("========== MYCO OUTPUT ==========")
        print(outputs)
        print("=================================")
        generated = outputs[0].get("generated_text", "")
        if isinstance(generated, list):
            last = generated[-1]
            return (last.get("content") if isinstance(last, dict) else str(last)).strip() or None
        return str(generated).strip() or None
    except Exception as exc:
        import traceback
        print(f"[Myco] Inference error: {exc}")
        traceback.print_exc()
        return None
        

def _llm_with_history(history: list, user_message: str, context: dict) -> str | None:
    """Call Gemma with full conversation history for the chat interface."""
    pipe = _get_pipeline()
    if not pipe:
        return None

    ctx = context or {}
    mushroom_line = ""
    if ctx.get("name"):
        poison_flag = " ⚠️ POISONOUS" if ctx.get("name") in POISONOUS else ""
        mushroom_line = (
            f"Current mushroom: {ctx['name']} ({ctx.get('rarity','?')}){poison_flag}. "
            f"Lore: {ctx.get('lore','?')}. "
            f"Edible: {ctx.get('edible','Unknown')}. Magic: {ctx.get('magic','Unknown')}."
        )

    system = (
        f"{SYSTEM_PROMPT}\n\n"
        f"{mushroom_line}\n"
        f"MycoDex entries: {ctx.get('collection_count', 0)}. "
        f"Mystery: {ctx.get('mystery_title', 'The Wrong Memory')}. "
        f"Score: {ctx.get('score', 0)} spores. Health: {ctx.get('health', 3)}/3."
    )

    messages = [{"role": "system", "content": system}]

    # Add conversation history (last 6 turns)
    for entry in history[-6:]:
        role    = entry.get("role", "assistant")
        content = entry.get("content", "")
        if isinstance(content, str) and content.strip():
            messages.append({"role": role, "content": content})

    messages.append({"role": "user", "content": user_message})

    try:
        outputs = _run_pipeline(pipe, messages)
        print("========== MYCO OUTPUT ==========")
        print(outputs)
        print("=================================")
        generated = outputs[0].get("generated_text", "")
        if isinstance(generated, list):
            last = generated[-1]
            return (last.get("content") if isinstance(last, dict) else str(last)).strip() or None
        return str(generated).strip() or None
    except Exception as exc:
        import traceback
        print(f"[Myco] Inference error: {exc}")
        traceback.print_exc()
        return None


# ---------------------------------------------------------------------------
# Context builder
# ---------------------------------------------------------------------------

def _ctx(current: dict | None, collection: list) -> dict:
    count = len(collection)
    chapter = MYSTERY_CHAPTERS[0]
    for c in MYSTERY_CHAPTERS:
        if count >= c["threshold"]:
            chapter = c
    score = _score_collection(collection)
    health = _health(current, collection)
    ctx: dict = {
        "collection_count": count,
        "mystery_title": chapter["title"],
        "mystery_clue":  chapter["clue"],
        "score":  score,
        "health": health,
    }
    if current:
        ctx.update({
            "name":    current.get("name", ""),
            "rarity":  current.get("rarity", "Common"),
            "habitat": current.get("habitat", ""),
            "lore":    current.get("lore", ""),
            "edible":  current.get("edible", "Unknown"),
            "magic":   current.get("magic", "Unknown"),
            "danger":  current.get("danger", "Unknown"),
        })
    return ctx


# ---------------------------------------------------------------------------
# Fallbacks
# ---------------------------------------------------------------------------

def _fallback_discover(current: dict) -> str:
    name   = current.get("name", "something")
    rarity = current.get("rarity", "Common")
    poison = current.get("name", "") in POISONOUS
    if poison:
        return f"Wait — {name}! I've seen this before... something feels very wrong. Don't touch it yet."
    if rarity == "Legendary":
        return f"Oh! Oh! A {name}! The whole clearing just went silent. This is from the Elder Map!"
    if rarity == "Rare":
        return f"A {name}... I can feel it humming. Something rare is here — maybe magical."
    return f"A {name}! Found near {current.get('habitat','the forest')}. Let me sense it first."

def _fallback_pick(current: dict) -> str:
    if _is_poisonous(current):
        return "💀 That was poisonous! I tried to stop you... the forest goes dark."
    return f"Got {current.get('name','it')}! +{RARITY_SCORE.get(current.get('rarity','Common'),10)} spores!"

def _fallback_study(current: dict) -> str:
    return f"I studied it carefully. Magic field updated. The clue: {RARITY_CLUES.get(current.get('rarity','Common'), '')}"

def _fallback_collect(current: dict) -> str:
    return f"Added {current.get('name','it')} to the MycoDex! The pages feel warmer."

def _fallback_whisper(current: dict) -> str:
    return "I followed the whisper... and remembered a path I've never walked. The mystery deepens."

def _fallback_chat(current: dict | None) -> str:
    if current:
        return f"I feel something strange about {current.get('name','this')}... stay close to me."
    return "The forest is full of secrets. Move to a clearing and search — I'll watch for danger."


# ---------------------------------------------------------------------------
# Mushroom helpers
# ---------------------------------------------------------------------------

def _choose_mushroom(catalog=None):
    mushrooms = tuple(load_mushrooms() if catalog is None else catalog)
    weights = [RARITY_WEIGHTS.get(m.rarity, 12) for m in mushrooms]
    return random.choices(mushrooms, weights=weights, k=1)[0]

def _is_poisonous(current: dict) -> bool:
    return current.get("name", "") in POISONOUS or current.get("danger") == "Poisonous"

def _score_value(current: dict) -> int:
    return RARITY_SCORE.get(current.get("rarity", "Common"), 10)

def _score_collection(collection: list) -> int:
    total = 0
    for e in collection:
        if e.get("game_over") == "Yes":
            continue
        total += int(e.get("score_delta") or _score_value(e))
    return max(0, total)

def _health(current: dict | None, collection: list) -> int:
    if current and current.get("game_over") == "Yes":
        return 0
    deaths = sum(1 for e in collection if e.get("game_over") == "Yes")
    return max(0, PLAYER_HEALTH - deaths)

def _mystery_state(count: int) -> dict:
    chapter, next_ch = MYSTERY_CHAPTERS[0], None
    for c in MYSTERY_CHAPTERS:
        if count >= c["threshold"]:
            chapter = c
        elif next_ch is None:
            next_ch = c
    next_line = (
        f"{next_ch['threshold'] - count} discoveries until {next_ch['title']}."
        if next_ch else "The Impossible Bloom is near. Follow the whisper."
    )
    return {
        "mystery_title": chapter["title"],
        "mystery_clue":  chapter["clue"],
        "mystery_next":  next_line,
    }

def _story_event(count: int) -> dict:
    return FOREST_EVENTS[count % len(FOREST_EVENTS)]

def _build_current(mushroom, collection: list) -> dict:
    count = len(collection)
    current = mushroom.to_dict()
    current["poison"]       = "Yes" if mushroom.name in POISONOUS else "No"
    current["score_total"]  = str(_score_collection(collection))
    current["health"]       = str(_health(current, collection))
    current["score_delta"]  = "0"
    current["clue"]         = RARITY_CLUES.get(mushroom.rarity, RARITY_CLUES["Common"])
    if count == 0:
        current["clue"] = f"First clue: {mushroom.name} marks the beginning of the Spore Door trail."
    event = _story_event(count)
    current.update({
        "event_title":       event["title"],
        "event_emoji":       event["emoji"],
        "myco_mood":         event["mood"],
        "reward_text":       "Discover, then pick or collect.",
    })
    current.update(_mystery_state(count))
    return current

def _append(history: list, role: str, content: str) -> list:
    return [*history, {"role": role, "content": content}]

def _safe_history(h) -> list:
    return list(h or welcome_history())

def _safe_collection(c) -> list:
    return list(c or [])


# ---------------------------------------------------------------------------
# PUBLIC GAME ACTIONS
# ---------------------------------------------------------------------------

def discover_mushroom(collection=None, catalog=None):
    """Discover a new mushroom. LLM narrates the moment."""
    coll    = _safe_collection(collection)
    mushroom = _choose_mushroom(catalog)
    current  = _build_current(mushroom, coll)
    ctx      = _ctx(current, coll)

    prompt = (
        f"The player just discovered a {mushroom.rarity} mushroom called {mushroom.name} "
        f"near {mushroom.habitat}. "
        f"{'WARNING: this mushroom is POISONOUS! React with fear and urgency.' if mushroom.name in POISONOUS else ''}"
        f"{'This is LEGENDARY — react with awe and excitement!' if mushroom.rarity == 'Legendary' else ''}"
        f"{'This is RARE — react with wonder and curiosity.' if mushroom.rarity == 'Rare' else ''}"
        f"Forest event: {current['event_title']}. Myco mood: {current['myco_mood']}. "
        f"Mystery chapter: {current['mystery_title']}. "
        "React in character. Hint at what to do next (Study, Pick, Follow Whisper, or Collect)."
    )
    reply = _llm(prompt, ctx) or _fallback_discover(current)
    history = _append(welcome_history(), "assistant", reply)
    return mushroom, current, history


def myco_reply(message=None, history=None, current=None, collection=None, position=None):
    """Player chats with Myco. LLM responds in character with full context."""
    print("MYCO_REPLY CALLED, message:", repr(message))
    hist  = _safe_history(history)
    coll  = _safe_collection(collection)
    clean = (message or "").strip()
    if not clean:
        return "", hist

    ctx   = _ctx(current, coll)
    reply = _llm_with_history(hist, clean, ctx)
    if not reply:
        reply = _fallback_chat(current)

    return "", _append(hist, "user", clean) + [{"role": "assistant", "content": reply}]

companion_reply = myco_reply


def collect_current(current=None, collection=None, history=None):
    """Collect mushroom into MycoDex. LLM narrates the entry."""
    coll = _safe_collection(collection)
    hist = _safe_history(history)

    if current is None:
        return coll, _append(hist, "assistant", "We need to find a mushroom first!")

    if collection_contains(coll, current["name"]):
        return coll, _append(hist, "assistant", f"{current['name']} is already in the MycoDex!")

    score_delta   = _score_value(current)
    score_total   = _score_collection(coll) + score_delta
    collected     = {
        **current,
        "score_delta": str(score_delta),
        "score_total": str(score_total),
        "health":      str(_health(current, coll)),
        "reward_text": f"+{score_delta} spores",
    }
    updated_coll  = [*coll, collected]

    ctx   = _ctx(current, coll)
    prompt = (
        f"The player just added {current['name']} ({current.get('rarity','Common')}) to the MycoDex! "
        f"+{score_delta} spores. Total score: {score_total}. "
        f"MycoDex now has {len(updated_coll)} entries. "
        "Celebrate this moment. Add a small lore detail or mystery hint."
    )
    reply = _llm(prompt, ctx) or _fallback_collect(current)
    return updated_coll, _append(hist, "assistant", reply)


def pick_current(current=None, collection=None, history=None):
    """Pick mushroom as game item. Poison = game over. LLM narrates dramatically."""
    coll = _safe_collection(collection)
    hist = _safe_history(history)

    if current is None:
        return coll, None, _append(hist, "assistant", "Find a mushroom first before picking!")

    ctx = _ctx(current, coll)

    if _is_poisonous(current):
        score_total = max(0, _score_collection(coll) + POISON_PENALTY)
        game_over   = {
            **current,
            "danger":      "Poisonous",
            "game_over":   "Yes",
            "health":      "0",
            "score_delta": str(POISON_PENALTY),
            "score_total": str(score_total),
            "reward_text": f"Poison! {POISON_PENALTY} spores · Game Over",
        }
        prompt = (
            f"DRAMATIC MOMENT: The player picked {current['name']} which is POISONOUS! "
            f"Game Over! Score drops by 25 to {score_total}. Health → 0. "
            "React with shock, sadness, and a dramatic farewell. Make it memorable."
        )
        reply = _llm(prompt, ctx) or _fallback_pick(current)
        return coll, game_over, _append(hist, "assistant", f"💀 {reply}")

    score_delta = _score_value(current)
    score_total = _score_collection(coll) + score_delta
    picked = {
        **current,
        "picked":      "Yes",
        "danger":      "Safe",
        "score_delta": str(score_delta),
        "score_total": str(score_total),
        "health":      str(_health(current, coll)),
        "reward_text": f"+{score_delta} spores",
    }

    if collection_contains(coll, picked["name"]):
        return coll, picked, _append(hist, "assistant", f"{picked['name']} already picked!")

    updated_coll = [*coll, picked]
    prompt = (
        f"The player safely picked {current['name']} ({current.get('rarity','Common')})! "
        f"+{score_delta} spores. Total: {score_total}. "
        "Celebrate! Make it feel like a platformer power-up moment."
    )
    reply = _llm(prompt, ctx) or _fallback_pick(current)
    return updated_coll, picked, _append(hist, "assistant", f"🍄 {reply}")


def follow_whisper(current=None, collection=None, history=None):
    """Follow the forest whisper. LLM reveals mystery fragments."""
    coll = _safe_collection(collection)
    hist = _safe_history(history)

    if current is None:
        return None, _append(hist, "assistant",
            "Myco cups one ear. The forest only whispers near mushrooms — search a clearing first.")

    ctx = _ctx(current, coll)

    if _is_poisonous(current) and current.get("studied") != "Yes":
        game_over = {**current, "danger": "Poisonous", "game_over": "Yes",
                     "health": "0", "score_total": str(max(0, _score_collection(coll) + POISON_PENALTY))}
        prompt = (
            f"The player followed a whisper but it led to POISON from {current['name']}! Game Over! "
            "React with horror and a haunting mystery revelation."
        )
        reply = _llm(prompt, ctx) or "💀 The whisper belonged to poison... Myco screams."
        return game_over, _append(hist, "assistant", reply)

    mystery = _mystery_state(len(coll) + 1)
    prompt = (
        f"The player followed a forest whisper near {current.get('name','a mushroom')}. "
        f"Mystery chapter revealed: {mystery['mystery_title']}. Clue: {mystery['mystery_clue']}. "
        f"Reveal this mystery fragment dramatically. "
        "Make Myco gasp or tremble. Hint that the MycoDex is alive and regrowing the lost forest."
    )
    reply = _llm(prompt, ctx) or _fallback_whisper(current)
    revealed = {
        **current,
        **mystery,
        "whisper_followed": "Yes",
    }
    return revealed, _append(hist, "assistant", f"🌌 {reply}")


def study_current(current=None, history=None):
    """Study mushroom. LLM gives a careful field observation."""
    hist = _safe_history(history)

    if current is None:
        return None, _append(hist, "assistant", "Nothing to study yet — find a mushroom first!")

    magic_hints = {"Legendary": "Strong, mysterious aura", "Rare": "Faint magical trace", "Common": "Tiny spore shimmer"}
    studied = {**current, "magic": magic_hints.get(current.get("rarity", "Common"), "Unknown"), "studied": "Yes"}

    ctx    = _ctx(studied, [])
    prompt = (
        f"Myco carefully studies {current['name']} ({current.get('rarity','Common')} rarity). "
        f"Habitat: {current.get('habitat','?')}. Lore: {current.get('lore','?')}. "
        f"Magic field is now: {studied['magic']}. "
        "Share one careful field observation. Warn if it might be dangerous. "
        "Be scientific but stay in character as a curious mushroom companion."
    )
    reply = _llm(prompt, ctx) or _fallback_study(current)
    return studied, _append(hist, "assistant", f"🔍 {reply}")


def eat_current(current=None, history=None):
    """Myco blocks eating. LLM reacts with personality."""
    hist = _safe_history(history)

    if current is None:
        return _append(hist, "assistant", "Find a mushroom before making snack decisions!")

    ctx    = _ctx(current, [])
    prompt = (
        f"The player wants to EAT {current.get('name','the mushroom')} "
        f"(edible status: {current.get('edible','Unknown')}). "
        "Block them! React with alarm or gentle scolding. Stay in character."
    )
    reply = _llm(prompt, ctx) or f"Myco blocks you — {current.get('name','this')} hasn't been studied enough!"
    return _append(hist, "assistant", reply)