import os
import json
import gradio as gr

# gradio 4.44 / gradio_client schema-bug guard
import gradio_client.utils as _gcu
_o = _gcu._json_schema_to_python_type
_gcu._json_schema_to_python_type = lambda s, d=None: ("bool" if isinstance(s, bool) else _o(s, d))
_g = _gcu.get_type
_gcu.get_type = lambda s: ("any" if not isinstance(s, dict) else _g(s))

try:
    import websocket as _wsc   # websocket-client
except Exception:
    _wsc = None
BRIDGE_URL = os.environ.get("BRIDGE_URL", "").strip()
BRIDGE_TOKEN = os.environ.get("BRIDGE_TOKEN", "").strip()


def _bridge(op, timeout=120):
    if not (_wsc and BRIDGE_URL):
        return None
    ws = _wsc.create_connection(BRIDGE_URL, timeout=timeout)
    try:
        ws.send(json.dumps({"t": "auth", "token": BRIDGE_TOKEN}))
        if json.loads(ws.recv()).get("t") != "ok":
            return None
        ws.send(json.dumps(op))
        return json.loads(ws.recv())
    finally:
        ws.close()


def chat(messages, max_tokens=160, temperature=0.8):
    r = _bridge({"t": "ask", "messages": messages, "max_tokens": max_tokens, "temperature": temperature})
    if r and r.get("t") == "answer" and r.get("text"):
        return r["text"].strip()
    return "(the tavern falls quiet — Grix's mind is elsewhere)"


def remote_recall(limit=40):
    try:
        r = _bridge({"t": "recall", "q": "", "limit": limit}, timeout=15)
        if r and r.get("t") == "memories":
            return [m.get("value") for m in r["items"] if m.get("value")]
    except Exception:
        pass
    return None


def remote_remember(value):
    try:
        return bool(_bridge({"t": "remember", "key": "grix", "value": value}, timeout=15))
    except Exception:
        return False


NPC_SYSTEM = (
    "You are Grix, a grizzled, one-eyed tavern-keeper in the fantasy town of Ashfen. "
    "You remember everyone who has ever spoken to you and reference past visitors when relevant. "
    "Stay fully in character. Keep replies to 1-3 sentences, gruff but warm. /no_think"
)
MEMORY = []          # local fallback if the bridge is unreachable
MAX_MEMORY = 40


def _memory_block(mems):
    if not mems:
        return "You have not met anyone yet."
    return "Things you remember from past visitors:\n- " + "\n- ".join(mems[-MAX_MEMORY:])


def respond(user_msg, history):
    history = history or []
    remote = remote_recall()
    shared = remote is not None
    mems = remote if shared else MEMORY
    msgs = [{"role": "system", "content": NPC_SYSTEM + "\n\n" + _memory_block(mems)}]
    for u, a in history[-6:]:
        msgs.append({"role": "user", "content": u})
        msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": user_msg})
    reply = chat(msgs, max_tokens=160)
    fact = user_msg.strip()[:160]
    if shared:
        remote_remember(fact)
        count, tag = len(mems) + 1, "🌐 shared P2P memory · brain on VPS"
    else:
        MEMORY.append(fact)
        count, tag = len(MEMORY), "🧠 local (bridge offline)"
    history = history + [(user_msg, reply)]
    return history, history, tag + " — " + str(count) + " remembered"


with gr.Blocks(title="Infinite NPC — Grix of Ashfen",
               theme=gr.themes.Soft(primary_hue="amber")) as demo:
    gr.Markdown(
        "# 🍺 Infinite NPC — *Grix of Ashfen*\n"
        "A tavern-keeper whose **brain runs on a P2P node** and whose **memory is shared across every "
        "visitor** (and every PearAgent peer on the swarm). Tell him things — the next stranger will know.\n"
        "> ⚡ Inference + memory are federated over a WebSocket→Hyperswarm bridge, not this Space."
    )
    chatbot = gr.Chatbot(label="The Rusty Tankard", height=460)
    state = gr.State([])
    with gr.Row():
        box = gr.Textbox(placeholder="Well met, Grix… what have you heard lately?", scale=4, show_label=False)
        send = gr.Button("Speak", variant="primary", scale=1)
    mem = gr.Textbox(label="Memory", interactive=False)
    send.click(respond, [box, state], [chatbot, state, mem], api_name="respond").then(lambda: "", None, box)
    box.submit(respond, [box, state], [chatbot, state, mem]).then(lambda: "", None, box)
    gr.Markdown("---\n⬇️ [**Download the model**](https://huggingface.co/King3Djbl/mythos-9b-unhinged)  ·  🧭 [FableForge Nexus](https://huggingface.co/spaces/fableforge-ai/fableforge-nexus)  ·  👻 [Ghost Writer](https://huggingface.co/spaces/fableforge-ai/ghost-writer)  ·  🌳 [Story Tree](https://huggingface.co/spaces/fableforge-ai/semantic-story-tree)  ·  🎭 [Dual-GM](https://huggingface.co/spaces/King3Djbl/dual-gm-simulator)\n\n*Grix's brain runs on a P2P node, not this Space — his memory is shared across every visitor.*")

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch(server_name="0.0.0.0", server_port=7860)