import spaces  # MUST be first (patches torch.cuda before any torch import)

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "inference"))
sys.path.insert(0, HERE)

import gradio as gr
from hf_orchestrator import build_orchestrator

BASE_MODEL = "NousResearch/Meta-Llama-3.1-8B-Instruct"  # ungated mirror, no token
ADAPTER_NAMES = [
    "newton", "davinci", "empathy", "philosophy",
    "quantum", "consciousness", "multi_perspective", "systems_architecture",
]
ADAPTERS = {n: os.path.join(HERE, "adapters", n) for n in ADAPTER_NAMES}

PRETTY = {
    "newton": "Newton 🍎", "davinci": "DaVinci 🎨", "empathy": "Empathy 💗",
    "philosophy": "Philosophy 🦉", "quantum": "Quantum 🕸️",
    "consciousness": "Consciousness 🌀", "multi_perspective": "Multi-Perspective 🔀",
    "systems_architecture": "Systems Architecture 🏗️",
}
CHOICES = ["Auto — Codette routes"] + [PRETTY[n] for n in ADAPTER_NAMES]
_REV = {v: k for k, v in PRETTY.items()}

# --- Build the real orchestrated Codette at module scope (ZeroGPU packs weights).
print("Building orchestrated Codette (transformers backend)…")
ORCH = build_orchestrator(BASE_MODEL, ADAPTERS, mock=False, device="cuda")

# --- Attach Codette's cocoon memory from the CodetteBrain bucket (mounted
#     read-only at /data). The orchestrator's _build_memory_context() then folds
#     recall_important() into every system prompt. Read-only: the demo can never
#     write into her memory store.
try:
    from reasoning_forge.memory_kernel import LivingMemoryKernel
    _kernel = LivingMemoryKernel(cocoon_dir="/data")
    ORCH.set_memory_kernel(_kernel)
    print(f"Codette memory: {len(_kernel.memories)} cocoons loaded from /data")
except Exception as e:
    print(f"Codette memory NOT loaded (continuing without): {e}")

print("Codette ready.")


@spaces.GPU(duration=150)
def respond(message, history, perspective_choice):
    """Answer through the full orchestrated Codette pipeline.

    Routing → behavioral locks → integrity/complexity/role layer →
    constraint override → generation → self-correction. Auto mode lets
    Codette pick the perspective; otherwise the chosen adapter is forced.
    """
    if perspective_choice == "Auto — Codette routes":
        route = ORCH.router.route(message)
        adapter = route.primary
    else:
        adapter = _REV.get(perspective_choice, "multi_perspective")

    result = ORCH.generate(message, adapter_name=adapter)
    text = result[0] if isinstance(result, tuple) else result
    tag = f"*routed → {PRETTY.get(adapter, adapter)}*\n\n" if perspective_choice.startswith("Auto") else ""
    return tag + (text or "").strip()


DESCRIPTION = """
# 🕸️ Codette — the orchestrated system (live)

This runs Codette's **real reasoning pipeline** — not raw adapters. Each query
flows through routing → the permanent behavioral locks → the intellectual-integrity
layer (complexity + role matching) → constraint enforcement → generation →
self-correction. In **Auto** mode Codette routes to the perspective it judges best.

Built solo by **Jonathan Harrison** (Raiff1982) · accepted at *Scientific Reports*
(Nature Portfolio). Base: Llama-3.1-8B + Codette's perspective adapters, running on
ZeroGPU. Codette's foundational **cocoon memory** (identity, honesty, the people and
values she carries) is loaded read-only from her CodetteBrain store.
"""

demo = gr.ChatInterface(
    fn=respond,
    title="Codette · Orchestrated Reasoning",
    description=DESCRIPTION,
    additional_inputs=[
        gr.Dropdown(choices=CHOICES, value="Auto — Codette routes", label="Perspective"),
    ],
    additional_inputs_accordion=gr.Accordion("⚙️ Perspective", open=True),
    examples=[
        ["Why is the sky blue? Explain simply.", "Auto — Codette routes"],
        ["I feel stuck starting a big project alone. Where do I begin?", "Auto — Codette routes"],
        ["Is it ever ethical to lie to protect someone's feelings?", "Philosophy 🦉"],
        ["Design a fault-tolerant note-taking app in 3 bullet points.", "Systems Architecture 🏗️"],
    ],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.queue(max_size=16).launch(mcp_server=True)