import os
import re

# ---- Fix gradio 4.44 / gradio_client schema bug (bool not iterable) ----
import gradio_client.utils as _gcu
_orig_j2p = _gcu._json_schema_to_python_type
def _safe_j2p(schema, defs=None):
    if isinstance(schema, bool):
        return "bool"
    return _orig_j2p(schema, defs)
_gcu._json_schema_to_python_type = _safe_j2p
_orig_gt = _gcu.get_type
def _safe_gt(schema):
    if not isinstance(schema, dict):
        return "any"
    return _orig_gt(schema)
_gcu.get_type = _safe_gt

import gradio as gr
from llama_cpp import Llama

MODEL_REPO = "King3Djbl/mythos-9b-unhinged"
GGUF_FILE = "mythos-9b-unhinged-Q4_K_M.gguf"

print("Loading", MODEL_REPO, GGUF_FILE)
llm = Llama.from_pretrained(
    repo_id=MODEL_REPO,
    filename=GGUF_FILE,
    n_ctx=2048,
    n_threads=os.cpu_count(),
    chat_format="chatml",
    verbose=False,
)
print("Model loaded.")


def chat(messages, max_tokens=220, temperature=0.9):
    out = llm.create_chat_completion(
        messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=0.9,
    )
    return out["choices"][0]["message"]["content"].strip()

FOOTER = '---\n### 🌌 More FableForge demos\n🍺 [Infinite NPC](https://huggingface.co/spaces/fableforge-ai/infinite-npc) · 👻 [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) · 🧭 [**Nexus hub →**](https://huggingface.co/spaces/fableforge-ai/fableforge-nexus)\n\n⬇️ [**Download the model**](https://huggingface.co/King3Djbl/mythos-9b-unhinged)  ·  🐦 [Share on X](https://twitter.com/intent/tweet?text=%F0%9F%94%AE%20FableForge%3A%20a%20galaxy%20of%20live%20AI%20demos%20powered%20by%20one%20open%209B%20model.%20Play%20them%20free%3A&url=https%3A//huggingface.co/spaces/fableforge-ai/fableforge-nexus)  ·  👽 [Post to Reddit](https://www.reddit.com/submit?url=https%3A//huggingface.co/spaces/fableforge-ai/fableforge-nexus&title=%F0%9F%94%AE%20FableForge%3A%20a%20galaxy%20of%20live%20AI%20demos%20powered%20by%20one%20open%209B%20model.%20Play%20them%20free%3A)\n\n*One open model. Infinite worlds. ⭐ Like this Space to boost it.*'

STYLES = ["Match my style", "Literary", "Hardboiled noir", "Whimsical fairytale", "Cosmic horror", "Punchy thriller"]


def ghost(draft, style, length):
    draft = draft or ""
    sys = "You are a ghost co-writer. Continue the user's text SEAMLESSLY from exactly where it stops. "
    sys += "Do not repeat their text, do not add commentary — output only the continuation. "
    if style and style != "Match my style":
        sys += "Write in this style: " + style + ". "
    cont = chat(
        [{"role": "system", "content": sys},
         {"role": "user", "content": draft if draft.strip() else "Begin a story."}],
        max_tokens=int(length), temperature=0.9,
    )
    joiner = "" if (not draft or draft.endswith((" ", "\n"))) else " "
    return (draft + joiner + cont).strip()


with gr.Blocks(title="Ghost Writer", theme=gr.themes.Soft(primary_hue="slate")) as demo:
    gr.Markdown(
        "# 👻 Ghost Writer\n"
        "Write a little, then let the ghost continue your prose. Keep hitting **Continue** to co-write.\n"
        "> ⏳ Free CPU — each continuation takes ~15-30s."
    )
    text = gr.Textbox(label="Your manuscript", lines=16,
                      placeholder="The lighthouse keeper hadn't spoken to another soul in three winters, until…")
    with gr.Row():
        style = gr.Dropdown(STYLES, value="Match my style", label="Voice", scale=2)
        length = gr.Slider(40, 300, value=120, step=20, label="Continue by (tokens)", scale=2)
        go = gr.Button("👻 Continue", variant="primary", scale=1)
    go.click(ghost, [text, style, length], text)
    gr.Markdown(FOOTER)

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