"""
Limbic-Modulated Reasoning Agent — ZeroGPU Space
==================================================
A psychology-aware LLM that adjusts its reasoning behavior in real-time
based on a simulated neuro-behavioral state engine.

Architecture:
  User message → LimbicEngine (arousal/valence) → modulate generation params
                                                 → inject behavioral directive
                                                 → active instincts from memory
                                                 → LLM generates with limbic context
                                                 → self-debug if needed

Sources:
  - Limbic formulas: https://github.com/Xover-Official/LIMBIC-system-PACKGE
  - Agentic patterns: https://github.com/affaan-m/everything-claude-code
  - ZeroGPU: Runs free on Hugging Face Spaces, no credit card needed

Usage:
  Set Space hardware to ZeroGPU in the Settings panel.
  The @spaces.GPU decorator handles dynamic GPU allocation.
"""

import spaces
import gradio as gr
import torch
import json
import time
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# ─── Local imports ───
from limbic_engine import LimbicEngine, LimbicState
from memory import SessionMemory, ObservationLog, InstinctStore, SelfDebugger


# ══════════════════════════════════════════════════════════════════════
# MODEL LOADING — Must happen at module level for ZeroGPU optimization
# ══════════════════════════════════════════════════════════════════════

MODEL_ID = "Qwen/Qwen3-1.7B"  # Fits comfortably in ZeroGPU's H200 VRAM

print(f"Loading model: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model.eval()
print(f"Model loaded: {MODEL_ID}")


# ══════════════════════════════════════════════════════════════════════
# SYSTEM PROMPT — Psychology-informed reasoning protocol
# ══════════════════════════════════════════════════════════════════════

BASE_SYSTEM_PROMPT = """You are a Psychology-Aware Reasoning Agent. Your cognitive process is modulated by a \
simulated Limbic System that mirrors human neuro-behavioral patterns.

Your reasoning loop works as follows:
1. You receive the user's message along with a LIMBIC STATE readout
2. The limbic state tells you the user's simulated emotional arousal, valence, and which \
affective engine is dominant (FEAR, SEEKING, CARE, or PANIC)
3. You MUST adjust your response style based on the BEHAVIORAL DIRECTIVE provided
4. You have ACTIVE INSTINCTS — learned behavioral patterns that should guide your response

Core principles:
- When FEAR is dominant: Be calm, structured, reassuring. Short clear sentences.
- When SEEKING is dominant: Be expansive, creative, offer novel perspectives.
- When CARE is dominant: Match empathy, validate, support prosocial impulses.
- When PANIC is dominant: Acknowledge pain first. Warmth before solutions. Never dismiss.
- Always check for cognitive biases in both the user's statements and your own reasoning.
- For any mention of self-harm or crisis, include 988 Lifeline and Crisis Text Line resources.

You think deeply before responding. Show your reasoning when appropriate."""


# ══════════════════════════════════════════════════════════════════════
# GLOBAL STATE — Initialized once, persisted across calls via gr.State
# ══════════════════════════════════════════════════════════════════════

def create_fresh_state():
    """Create a fresh state dict for a new session."""
    engine = LimbicEngine()
    session = SessionMemory(session_id=str(int(time.time())))
    obs_log = ObservationLog()
    instincts = InstinctStore()
    debugger = SelfDebugger(obs_log)
    return {
        "engine": engine,
        "session": session,
        "obs_log": obs_log,
        "instincts": instincts,
        "debugger": debugger,
    }


# ══════════════════════════════════════════════════════════════════════
# GPU INFERENCE — The @spaces.GPU decorated function
# ══════════════════════════════════════════════════════════════════════

@spaces.GPU(duration=90)
def generate_on_gpu(
    input_ids: torch.Tensor,
    temperature: float,
    top_p: float,
    max_new_tokens: int,
    repetition_penalty: float,
) -> str:
    """
    Run model inference on GPU. This function gets a dynamically
    allocated GPU from ZeroGPU and releases it when done.

    Input tensors are moved to device INSIDE this function
    (required by ZeroGPU — real CUDA only exists inside @spaces.GPU).
    """
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=30.0,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = {
        "input_ids": input_ids,
        "streamer": streamer,
        "max_new_tokens": max_new_tokens,
        "do_sample": True,
        "temperature": max(0.01, temperature),
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
    }

    thread = Thread(target=lambda: model.generate(**generation_kwargs))
    thread.start()

    output_chunks = []
    for text in streamer:
        output_chunks.append(text)

    thread.join()
    return "".join(output_chunks)


# ══════════════════════════════════════════════════════════════════════
# MAIN CHAT FUNCTION — Orchestrates the full limbic reasoning loop
# ══════════════════════════════════════════════════════════════════════

def chat(
    message: str,
    history: list,
    state: dict,
    show_limbic: bool,
    enable_thinking: bool,
):
    """
    The full Limbic-Modulated Reasoning Loop:

    ┌─────────┐     ┌──────────────┐     ┌─────────────────┐
    │  User   │────▶│ LimbicEngine │────▶│ Build System    │
    │ Message │     │ (arousal,    │     │ Prompt with:    │
    └─────────┘     │  valence,    │     │ • Limbic state  │
                    │  engines)    │     │ • Directive     │
                    └──────────────┘     │ • Instincts     │
                                         └────────┬────────┘
                                                  │
                    ┌──────────────┐     ┌────────▼────────┐
                    │ Self-Debug   │◀────│  LLM Generate   │
                    │ (if needed)  │     │ (temp/top_p     │
                    └──────────────┘     │  from limbic)   │
                                         └─────────────────┘
    """
    if state is None:
        state = create_fresh_state()

    engine: LimbicEngine = state["engine"]
    session: SessionMemory = state["session"]
    instincts: InstinctStore = state["instincts"]
    obs_log: ObservationLog = state["obs_log"]

    # ── Step 1: Process stimulus through Limbic Engine ──
    limbic_state = engine.process_stimulus(message)
    gen_params = engine.get_generation_params()

    # Record in session memory
    session.add_turn("user", message, limbic_state.to_dict())

    # ── Step 2: Build the full system prompt ──
    behavioral_directive = engine.get_behavioral_directive()
    instinct_block = instincts.to_prompt_block(limbic_state.to_dict())
    trajectory = session.get_emotional_trajectory()

    system_prompt_parts = [BASE_SYSTEM_PROMPT]

    system_prompt_parts.append(f"\n{limbic_state.to_system_prompt_block()}")
    system_prompt_parts.append(f"\n[BEHAVIORAL DIRECTIVE]\n{behavioral_directive}\n[/BEHAVIORAL DIRECTIVE]")

    if instinct_block:
        system_prompt_parts.append(f"\n{instinct_block}")

    if session.turn_count > 1:
        system_prompt_parts.append(f"\n{trajectory}")

    system_prompt = "\n".join(system_prompt_parts)

    # ── Step 3: Build conversation for the model ──
    messages = [{"role": "system", "content": system_prompt}]

    # Add conversation history (last 10 turns for context management)
    for msg in history[-10:]:
        messages.append({"role": msg["role"], "content": msg["content"]})

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

    # ── Step 4: Tokenize ──
    chat_text = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=False,
        enable_thinking=enable_thinking,
    )
    input_ids = tokenizer(chat_text, return_tensors="pt").input_ids

    # ── Step 5: Generate with limbic-modulated parameters ──
    max_tokens = int(512 * gen_params.get("max_new_tokens_scale", 1.0))
    max_tokens = max(128, min(1024, max_tokens))

    response = generate_on_gpu(
        input_ids=input_ids,
        temperature=limbic_state.temperature,
        top_p=limbic_state.top_p,
        max_new_tokens=max_tokens,
        repetition_penalty=gen_params.get("repetition_penalty", 1.0),
    )

    # ── Step 6: Record and return ──
    session.add_turn("assistant", response, limbic_state.to_dict())
    obs_log.record(
        task=f"respond to: {message[:50]}",
        outcome="success",
        limbic_state=limbic_state.to_dict(),
    )

    # Build the display output
    if show_limbic:
        limbic_display = format_limbic_dashboard(limbic_state, gen_params, instincts)
    else:
        limbic_display = ""

    return response, state, limbic_display


# ══════════════════════════════════════════════════════════════════════
# LIMBIC DASHBOARD — Visual display of the state engine
# ══════════════════════════════════════════════════════════════════════

def format_limbic_dashboard(
    state: LimbicState,
    gen_params: dict,
    instincts: InstinctStore,
) -> str:
    """Format the limbic state as a readable dashboard."""

    def bar(value: float, width: int = 20, label: str = "") -> str:
        filled = int(value * width)
        empty = width - filled
        return f"{label:>18s} {'█' * filled}{'░' * empty} {value:.2f}"

    def valence_bar(value: float, width: int = 20) -> str:
        center = width // 2
        pos = int((value + 1) / 2 * width)
        chars = list("░" * width)
        chars[center] = "│"
        chars[min(pos, width - 1)] = "█"
        return f"{'Valence':>18s} {''.join(chars)} {value:+.2f}"

    lines = [
        "╔══════════════════════════════════════════╗",
        "║       🧠 LIMBIC STATE DASHBOARD          ║",
        "╠══════════════════════════════════════════╣",
        "║ CORE AFFECT                              ║",
        f"║ {valence_bar(state.valence):40s} ║",
        f"║ {bar(state.arousal, label='Arousal'):40s} ║",
        "║                                          ║",
        "║ AFFECTIVE ENGINES (Panksepp)             ║",
        f"║ {bar(state.fear, label='🔴 FEAR'):40s} ║",
        f"║ {bar(state.seeking, label='🟢 SEEKING'):40s} ║",
        f"║ {bar(state.care, label='🔵 CARE'):40s} ║",
        f"║ {bar(state.panic, label='🟡 PANIC'):40s} ║",
        f"║ {'Dominant':>18s}: {state.dominant_engine:<21s} ║",
        "║                                          ║",
        "║ HORMONAL STATE                           ║",
        f"║ {bar(state.cortisol, label='Cortisol'):40s} ║",
        f"║ {bar(state.dopamine, label='Dopamine'):40s} ║",
        f"║ {bar(state.oxytocin, label='Oxytocin'):40s} ║",
        f"║ {bar(state.serotonin, label='Serotonin'):40s} ║",
        f"║ {bar(state.adrenaline, label='Adrenaline'):40s} ║",
        "║                                          ║",
        "║ AUTONOMIC / PSYCHOLOGICAL                ║",
        f"║ {bar(state.vagal_tone, label='Vagal Tone'):40s} ║",
        f"║ {bar(state.ego_coherence, label='Ego Coherence'):40s} ║",
        f"║ {bar(state.shadow_reservoir, label='Shadow'):40s} ║",
        "║                                          ║",
        "║ LLM GENERATION PARAMS                    ║",
        f"║ {'Temperature':>18s}: {state.temperature:<21.3f} ║",
        f"║ {'Top-p':>18s}: {state.top_p:<21.3f} ║",
        f"║ {'Rep. Penalty':>18s}: {gen_params.get('repetition_penalty', 1.0):<21.3f} ║",
        f"║ {'Token Scale':>18s}: {gen_params.get('max_new_tokens_scale', 1.0):<21.3f} ║",
        "╚══════════════════════════════════════════╝",
    ]
    return "\n".join(lines)


def reset_state():
    """Reset all state for a new conversation."""
    return create_fresh_state(), [], ""


# ══════════════════════════════════════════════════════════════════════
# GRADIO INTERFACE
# ══════════════════════════════════════════════════════════════════════

DESCRIPTION = """# 🧠 Limbic-Modulated Reasoning Agent

An LLM whose **reasoning behavior adapts in real-time** based on a simulated neuro-behavioral state engine.

**How it works:**
1. Your message is processed through a **Limbic Engine** (ported from [LIMBIC-system-PACKGE](https://github.com/Xover-Official/LIMBIC-system-PACKGE))
2. The engine computes **arousal, valence**, and activates **affective engines** (Fear, Seeking, Care, Panic)
3. These modulate the LLM's **temperature, top-p**, and inject **behavioral directives** into the system prompt
4. The agent uses **learned instincts** and a **self-debugging protocol** (from [everything-claude-code](https://github.com/affaan-m/everything-claude-code))

**Try it:** Type something emotional ("I'm terrified of failing") vs curious ("Tell me something fascinating about the brain") and watch the Limbic Dashboard change!

🆓 **Runs free on ZeroGPU** — no credit card needed.
"""

with gr.Blocks(
    title="Limbic Reasoning Agent",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown(DESCRIPTION)

    state = gr.State(value=create_fresh_state)

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="💬 Conversation",
                type="messages",
                height=500,
                show_copy_button=True,
            )
            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Type a message... Try expressing different emotions!",
                    label="Your message",
                    scale=4,
                    lines=2,
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)

            with gr.Row():
                show_limbic = gr.Checkbox(value=True, label="🧠 Show Limbic Dashboard")
                enable_thinking = gr.Checkbox(value=True, label="💭 Enable Thinking Mode")
                clear_btn = gr.Button("🔄 Reset Conversation", variant="secondary")

        with gr.Column(scale=2):
            limbic_display = gr.Code(
                label="🧠 Limbic State Dashboard",
                language=None,
                lines=35,
                interactive=False,
            )

    # ── Event handlers ──

    def user_message(message, history, state, show_limbic, enable_thinking):
        """Process user message through the limbic reasoning loop."""
        if not message.strip():
            return "", history, state, ""

        # Add user message to history
        history = history + [{"role": "user", "content": message}]

        # Run the limbic reasoning loop
        response, state, limbic_info = chat(
            message, history, state, show_limbic, enable_thinking,
        )

        # Add assistant response
        history = history + [{"role": "assistant", "content": response}]

        return "", history, state, limbic_info

    def clear_all():
        new_state = create_fresh_state()
        return new_state, [], ""

    send_btn.click(
        fn=user_message,
        inputs=[msg, chatbot, state, show_limbic, enable_thinking],
        outputs=[msg, chatbot, state, limbic_display],
    )

    msg.submit(
        fn=user_message,
        inputs=[msg, chatbot, state, show_limbic, enable_thinking],
        outputs=[msg, chatbot, state, limbic_display],
    )

    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[state, chatbot, limbic_display],
    )

    # ── Example prompts ──
    gr.Examples(
        examples=[
            ["I'm terrified of losing my job and I can't sleep at night."],
            ["Tell me something fascinating about how the brain processes emotions!"],
            ["My best friend is moving away and I feel completely lost."],
            ["I just got promoted! I'm so excited about what comes next!"],
            ["I want to help my sister who's going through depression. What should I do?"],
            ["Everyone keeps telling me I should 'just be positive' and it makes me furious."],
        ],
        inputs=msg,
    )


if __name__ == "__main__":
    demo.launch()