#!/usr/bin/env python3
"""
Phox on Hugging Face Spaces — BQSM Wave-Rider Inference Engine.

This app wraps the Phoenix Brain C inference binary in a Gradio interface.
At Space startup it:
  1. Downloads the tokenizer from HF Hub (tokenizers lib, no PyTorch)
  2. On button click: downloads pre-built phoenix binary + model from HF Hub
  3. Starts phoenix in --chat mode with shared-memory ring buffers
  4. Provides a chat UI via Gradio

The phoenix binary runs on CPU (AVX2) — the wave-rider physics engine
doesn't use PyTorch, so ZeroGPU doesn't accelerate it directly.
But the Space provides a free public endpoint to chat with Phox.
"""
import os
import sys
import time
import json
import mmap
import struct
import subprocess
import threading

import gradio as gr
from huggingface_hub import hf_hub_download

# ── Config ──
SPACE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_REPO = "compunerd/emerging-systems-models"
MODEL_FILENAME = "gemma4-12b-ternary-normed.bqsm"
MODEL_SUBDIR = "/tmp/phoenix_models"
MODEL_PATH = os.path.join(MODEL_SUBDIR, MODEL_FILENAME)
STATE_FILE = "/tmp/phoenix_state.jsonl"
STATE_LOG = "/tmp/phoenix_daemon.log"
RING_IN_PATH = "/tmp/phoenix_ring_in"
RING_OUT_PATH = "/tmp/phoenix_ring_out"
RING_CAPACITY = 4096
RING_BUF_SIZE = RING_CAPACITY * 4 + 16
PHOENIX_BIN = "/tmp/phoenix"
MIXER_BIN = "/tmp/mixer"
CHAT_TIMEOUT_S = 120
SENTINEL_QUIT = 0xFFFFFFFE

# ── Global state ──
_phoenix_proc = None
_mixer_proc = None
_ring_in_mm = None
_ring_out_mm = None
_ring_in_fd = None
_ring_out_fd = None
_state_lock = threading.Lock()
_chat_messages = []
_boot_status = {"ready": False, "phase": "idle", "message": "Not started"}
_boot_thread = None
_tokenizer = None


def _update_boot(phase, message, ready=False):
    """Update boot status (thread-safe)."""
    with _state_lock:
        _boot_status["phase"] = phase
        _boot_status["message"] = message
        _boot_status["ready"] = ready


def load_tokenizer():
    """Download and load the Gemma 4 tokenizer from HF Hub."""
    global _tokenizer
    print("[Phox Space] Loading Gemma 4 tokenizer from Hub...")
    try:
        tok_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename="tokenizer/tokenizer.json")
        from tokenizers import Tokenizer
        _tokenizer = Tokenizer.from_file(tok_path)
        # Build a simple vocab-size attribute for compatibility
        _tokenizer.vocab_size = len(_tokenizer.get_vocab())
        print(f"[Phox Space] Tokenizer loaded (vocab={_tokenizer.vocab_size})")
    except Exception as e:
        print(f"[Phox Space] Tokenizer download failed: {e}")
        # Fallback: try local paths
        for tok_path in [
            os.path.join(MODEL_SUBDIR, "tokenizer", "tokenizer.json"),
            "/home/compunerd/models/gemma4-tokenizer/tokenizer.json",
        ]:
            try:
                from tokenizers import Tokenizer
                _tokenizer = Tokenizer.from_file(tok_path)
                _tokenizer.vocab_size = len(_tokenizer.get_vocab())
                print(f"[Phox Space] Tokenizer loaded from {tok_path}")
                break
            except Exception:
                pass
        if _tokenizer is None:
            print("[Phox Space] WARNING: No tokenizer available — chat will not work")


def boot_phoenix():
    """Full startup sequence: download model, build binary, start phoenix."""
    # Phase 0: Tokenizer
    _update_boot("tokenizer", "Loading tokenizer from Hub...")
    load_tokenizer()

    # Phase 1: Download model
    _update_boot("download", "Downloading BQSM model from HF Hub...")
    try:
        os.makedirs(MODEL_SUBDIR, exist_ok=True)
        if not os.path.exists(MODEL_PATH):
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=MODEL_FILENAME,
                local_dir=MODEL_SUBDIR,
                local_dir_use_symlinks=False,
            )
        _update_boot("download", f"Model ready: {os.path.getsize(MODEL_PATH)} bytes")
    except Exception as e:
        _update_boot("download", f"Model download FAILED: {e}")
        return

    # Phase 2: Download pre-built binary from HF Hub (avoids compilation on Space)
    _update_boot("build", "Checking for phoenix binary...")
    # Only download if binary doesn't exist (cached after first download)
    if not os.path.exists(PHOENIX_BIN):
        try:
            _update_boot("build", "Downloading phoenix binary from Hub...")
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename="phoenix",
                local_dir="/tmp",
                local_dir_use_symlinks=False,
                force_filename="phoenix",
            )
            os.chmod(PHOENIX_BIN, 0o755)
        except Exception as e:
            # Fallback: compile from source
            _update_boot("build", f"Binary download failed, compiling... ({e})")
            src_path = os.path.join(SPACE_DIR, "phoenix_brain.c")
            if os.path.exists(src_path):
                cc = os.environ.get("CC", "cc")
                cmd = [cc, "-O3", "-std=c11", "-march=native", "-fopenmp",
                       src_path, "-o", PHOENIX_BIN, "-lm"]
                result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
                if result.returncode != 0:
                    _update_boot("build", f"Build FAILED: {result.stderr[:500]}")
                    return
            else:
                _update_boot("build", "No binary or source found!")
                return
    _update_boot("build", "Binary ready")

    # Phase 3: Start phoenix
    _update_boot("start", "Starting phoenix in chat mode...")
    global _phoenix_proc
    with _state_lock:
        if _phoenix_proc and _phoenix_proc.poll() is None:
            _update_boot("start", "Already running", ready=True)
            return

        init_ring_buffers()
        _phoenix_proc = subprocess.Popen(
            [PHOENIX_BIN, MODEL_PATH, "--chat"],
            stdout=open(STATE_LOG, 'a'),
            stderr=subprocess.STDOUT,
        )
        _update_boot("start", f"Phoenix PID {_phoenix_proc.pid} — loading model (~40s)...")
        # Start a monitor thread to check when phoenix enters chat mode
        threading.Thread(target=_monitor_phoenix, daemon=True).start()

def _monitor_phoenix():
    """Background thread: watches daemon log for 'Chat Mode', then sets ready."""
    import time as _time
    for _ in range(120):  # Check for up to 2 minutes
        _time.sleep(0.5)
        try:
            size = os.path.getsize(STATE_LOG)
            with open(STATE_LOG, 'rb') as f:
                f.seek(max(0, size - 5000))
                log = f.read().decode('utf-8', errors='replace')
            if "Chat Mode" in log or "ring buffer" in log.lower() or "ready" in log.lower():
                _update_boot("ready", "Phoenix is ready to chat!")
                return
            # Check if process died
            with _state_lock:
                if _phoenix_proc and _phoenix_proc.poll() is not None:
                    _update_boot("error", f"Phoenix exited (code {_phoenix_proc.returncode})")
                    return
        except (FileNotFoundError, PermissionError):
            pass
    _update_boot("ready", "Phoenix boot timed out — may still be initializing.")


def init_ring_buffers():
    """Create ring buffer files with initialized headers."""
    for path in [RING_IN_PATH, RING_OUT_PATH]:
        try:
            os.unlink(path)
        except FileNotFoundError:
            pass
        fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o666)
        os.write(fd, b'\x00' * RING_BUF_SIZE)
        os.lseek(fd, 0, 0)
        os.write(fd, struct.pack('<IIII', 0, 0, 0, RING_CAPACITY))
        os.close(fd)

    global _ring_in_fd, _ring_out_fd, _ring_in_mm, _ring_out_mm
    _ring_in_fd = os.open(RING_IN_PATH, os.O_RDWR)
    _ring_in_mm = mmap.mmap(_ring_in_fd, RING_BUF_SIZE,
                           mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)
    _ring_out_fd = os.open(RING_OUT_PATH, os.O_RDWR)
    _ring_out_mm = mmap.mmap(_ring_out_fd, RING_BUF_SIZE,
                            mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)


def ring_push(mm, val):
    head = struct.unpack_from('<I', mm, 0)[0]
    tail = struct.unpack_from('<I', mm, 4)[0]
    sz = struct.unpack_from('<I', mm, 8)[0]
    next_tail = (tail + 1) % RING_CAPACITY
    if next_tail == head:
        return False
    struct.pack_into('<I', mm, 16 + tail * 4, val)
    struct.pack_into('<I', mm, 4, next_tail)
    struct.pack_into('<I', mm, 8, sz + 1)
    return True


def ring_pop(mm):
    head = struct.unpack_from('<I', mm, 0)[0]
    tail = struct.unpack_from('<I', mm, 4)[0]
    sz = struct.unpack_from('<I', mm, 8)[0]
    if sz == 0:
        return None
    val = struct.unpack_from('<I', mm, 16 + head * 4)[0]
    next_head = (head + 1) % RING_CAPACITY
    struct.pack_into('<I', mm, 0, next_head)
    struct.pack_into('<I', mm, 8, sz - 1)
    return val


def send_tokens_to_phoenix(token_ids):
    if not _ring_in_mm:
        return 0
    pushed = 0
    for tid in token_ids:
        if ring_push(_ring_in_mm, tid):
            pushed += 1
    return pushed


def poll_predictions(timeout_s=CHAT_TIMEOUT_S):
    results = []
    start = time.time()
    while time.time() - start < timeout_s:
        tok = ring_pop(_ring_out_mm)
        if tok is not None:
            if tok == SENTINEL_QUIT:
                break
            results.append(tok)
        else:
            time.sleep(0.05)
    return results


def start_phoenix_manual():
    """Start phoenix: download model, build binary, launch in chat mode."""
    global _phoenix_proc
    with _state_lock:
        if _phoenix_proc and _phoenix_proc.poll() is None:
            return "Already running"
    # Run boot in a thread so the Gradio event loop isn't blocked
    threading.Thread(target=boot_phoenix, daemon=True).start()
    return " Booting... check Engine Status for progress"


def start_mixer():
    """Start the mixing ring engine (lighter-weight alternative to phoenix)."""
    global _mixer_proc
    # Check if already running
    if _mixer_proc and _mixer_proc.poll() is None:
        return "Mixer already running"
    # Download binary to /tmp
    if not os.path.exists(MIXER_BIN):
        try:
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename="mixer",
                local_dir="/tmp",
                local_dir_use_symlinks=False,
            )
            os.chmod(MIXER_BIN, 0o755)
        except Exception as e:
            return f"Failed to download mixer binary: {e}"
    # Ensure model exists
    if not os.path.exists(MODEL_PATH):
        os.makedirs(MODEL_SUBDIR, exist_ok=True)
        try:
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=MODEL_FILENAME,
                local_dir=MODEL_SUBDIR,
                local_dir_use_symlinks=False,
            )
        except Exception as e:
            return f"Failed to download model: {e}"
    # Launch mixer binary — it runs inference and outputs to stdout
    init_ring_buffers()
    _mixer_proc = subprocess.Popen(
        [MIXER_BIN, MODEL_PATH, "-s", "2"],
        stdout=open(STATE_LOG, 'a'),
        stderr=subprocess.STDOUT,
    )
    return f"Mixer started (PID {_mixer_proc.pid}) — check daemon log"


def stop_phoenix():
    global _phoenix_proc, _ring_in_mm, _ring_out_mm, _ring_in_fd, _ring_out_fd
    # Send quit signal via ring buffer (no lock needed for ring_push)
    if _ring_in_mm:
        try:
            ring_push(_ring_in_mm, SENTINEL_QUIT)
        except Exception:
            pass
    with _state_lock:
        if _phoenix_proc and _phoenix_proc.poll() is None:
            _phoenix_proc.terminate()
    # Close ring buffers outside the lock to avoid blocking
    if _ring_in_mm:
        try:
            _ring_in_mm.close()
        except Exception:
            pass
        _ring_in_mm = None
    if _ring_out_mm:
        try:
            _ring_out_mm.close()
        except Exception:
            pass
        _ring_out_mm = None
    if _ring_in_fd is not None:
        try:
            os.close(_ring_in_fd)
        except Exception:
            pass
        _ring_in_fd = None
    if _ring_out_fd is not None:
        try:
            os.close(_ring_out_fd)
        except Exception:
            pass
        _ring_out_fd = None
    return "Stopped"


def get_boot_status():
    """Return current boot status for the UI."""
    # Read without lock to avoid contention with boot_phoenix thread
    try:
        return dict(_boot_status)
    except Exception:
        return {"ready": False, "phase": "error", "message": "Status unavailable"}


def phox_chat(message, history):
    """Main chat function — writes tokens to ring buffer, reads predictions."""
    if not _ring_in_mm or not _ring_out_mm:
        return "Engine not initialized. Please wait for startup."

    with _state_lock:
        engine_running = (_phoenix_proc is not None and _phoenix_proc.poll() is None) or \
                         (_mixer_proc is not None and _mixer_proc.poll() is None)
    if not engine_running:
        return "Engine not running. Please click 'Start Engine' or 'Start Mixer' first."

    if _tokenizer is None:
        return "Tokenizer not loaded."

    try:
        token_ids = _tokenizer.encode(message).ids
    except Exception as e:
        return f"Tokenization error: {e}"

    if not token_ids:
        return "Empty input."

    pushed = send_tokens_to_phoenix(token_ids)
    if pushed == 0:
        return "Ring buffer full. Try again."

    pred_ids = poll_predictions()
    # Debug: log what we got
    import os
    with open("/tmp/chat_debug.log", "a") as f:
        f.write(f"Input tokens: {token_ids}, Output tokens: {pred_ids}\n")
    if not pred_ids:
        return "No prediction received (timeout). Try again."

    response = _tokenizer.decode(pred_ids, skip_special_tokens=True)
    return response.strip() if response else "[no output]"


# ── Gradio interface ──

with gr.Blocks(title="Phox — Wave-Rider Brain") as demo:
    gr.Markdown("""
    # 🌀 Phox — Wave-Rider Inference Engine
    
    **Two engines available:**
    - **Phox Brain** (heavy): 2.98GB model, ~1.1 tok/s, full wave-rider physics
    - **Mixing Ring** (fast): 100-ring topology, ~8000+ tok/s, parallel settle
    
    Both use the same Gemma-4-12B ternary weights via mmap. Click the appropriate
    button to start, then chat below.
    """)

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.ChatInterface(
                fn=phox_chat,
                title="Talk to Phox",
                description="Type a message and watch the oscillators work.",
            )
        with gr.Column(scale=1):
            gr.Markdown("### Engine Status")
            status_json = gr.JSON(label="Boot Status")
            gr.Markdown("### Controls")
            with gr.Row():
                start_btn = gr.Button("Start Engine", variant="primary")
                mixer_btn = gr.Button("Start Mixer (fast)", variant="secondary")
                stop_btn = gr.Button("Stop Engine", variant="secondary")
            status_text = gr.Textbox(label="Engine Output", interactive=False)

            start_btn.click(fn=start_phoenix_manual, outputs=status_text)
            mixer_btn.click(fn=start_mixer, outputs=status_text)
            stop_btn.click(fn=stop_phoenix, outputs=status_text)

    # Periodic status refresh using gr.Timer (Gradio 6.x API)
    def tick():
        base = f"Phase: {_boot_status['phase']} | {_boot_status['message']}"
        # Show last line of daemon log if available (read tail only)
        try:
            import os as _os
            if _os.path.exists(STATE_LOG):
                size = _os.path.getsize(STATE_LOG)
                with open(STATE_LOG, 'rb') as f:
                    f.seek(max(0, size - 500))
                    tail = f.read().decode('utf-8', errors='replace')
                lines = [l.strip() for l in tail.split('\n') if l.strip()]
                if lines:
                    last = lines[-1]
                    if last:
                        base += f"\n[daemon] {last[:200]}"
        except (FileNotFoundError, PermissionError):
            pass
        return base
    timer1 = gr.Timer(3, active=True)
    timer1.tick(fn=get_boot_status, outputs=status_json)
    timer2 = gr.Timer(2, active=True)
    timer2.tick(fn=tick, outputs=status_text)


if __name__ == "__main__":
    # Load tokenizer at startup only (fast, ~32MB)
    # Model download + binary build happen on button click
    _boot_thread = threading.Thread(target=load_tokenizer, daemon=True)
    _boot_thread.start()
    # Don't auto-boot phoenix — user clicks "Start Engine"
    _update_boot("idle", "Click 'Start Engine' to boot the BQSM wave-rider (mixer auto-downloads)")
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        share=False,
    )
# Force rebuild Sat Aug  8 02:08:01 PM EDT 2026