"""
HOLD BRIDGE - Connects CausationHold to Viewer
═══════════════════════════════════════════════

Wires the real session data to the glass box viewer.
No synthetic data - pulls directly from CausationHold state.
"""

import threading
from typing import Optional, Dict, Any, Callable
from dataclasses import asdict

from .session import CausationHold, InferenceStep
from .viewer import HoldViewer, HoldViewerConfig


class HoldBridge:
    """
    Bridge between CausationHold and HoldViewer.
    
    Automatically displays HOLD points in the browser viewer,
    using real data from the session.
    
    Usage:
        from cascade.hold import CausationHold, HoldBridge
        
        hold = CausationHold()
        bridge = HoldBridge(hold, port=8080)
        bridge.start()
        
        # Now all capture() calls will pause and show in browser
        session = hold.begin_session("my_agent")
        choice, feedback = hold.capture(...)  # Opens in browser
    """
    
    def __init__(
        self,
        hold: CausationHold,
        port: int = 8080,
        auto_open: bool = True,
        action_labels: Optional[Dict[int, Dict[str, str]]] = None,
    ):
        self.hold = hold
        self.viewer = HoldViewer(HoldViewerConfig(
            port=port,
            auto_open=auto_open,
            action_labels=action_labels or {},
        ))
        self._original_capture = None
        self._running = False
        
    def start(self):
        """Start the bridge - intercepts capture() calls."""
        if self._running:
            return
            
        self._running = True
        self.viewer.start()
        
        # Hook into the hold system
        self._original_capture = self.hold.capture
        self.hold.capture = self._bridged_capture
        
    def stop(self):
        """Stop the bridge."""
        if not self._running:
            return
            
        self._running = False
        self.viewer.stop()
        
        if self._original_capture:
            self.hold.capture = self._original_capture
            self._original_capture = None
            
    def _bridged_capture(
        self,
        input_context: Dict[str, Any],
        candidates: list,
        **kwargs,
    ) -> tuple:
        """
        Intercept capture and route to viewer.
        """
        # Get the step data from hold
        step = self.hold.current_step
        session = self.hold.current_session
        
        # Build viewer data from real session state
        probs = [c.get('probability', 0) for c in candidates]
        ai_choice = probs.index(max(probs)) if probs else 0
        
        # Extract imagination if available
        imagination = {}
        for i, c in enumerate(candidates):
            if 'trajectory' in c or 'expected_value' in c:
                imagination[i] = {
                    'trajectory': c.get('trajectory', []),
                    'expected_value': c.get('expected_value', 0),
                }
        
        hold_data = {
            'hold_count': session.step_count if session else 0,
            'merkle': step.merkle_root if step else None,
            'brain_id': self.hold._active_agent_id,
            'action_probs': probs,
            'ai_choice': ai_choice,
            'value': step.value_estimate if step else 0,
            'observation': input_context,
            'imagination': imagination,
            'attention': input_context.get('attention', {}),
            'reasoning': [c.get('reasoning', '') for c in candidates if c.get('reasoning')],
        }
        
        # Show in viewer and wait
        self.viewer.show_hold(hold_data)
        action = self.viewer.wait_for_resolution()
        
        # Now call original capture with the resolved action
        # (the original will handle the actual step recording)
        return self._original_capture(
            input_context=input_context,
            candidates=candidates,
            override_choice=action,
            **kwargs,
        )


def with_viewer(hold: CausationHold, port: int = 8080) -> CausationHold:
    """
    Decorator-style: wrap a CausationHold with browser viewer.
    
    Usage:
        hold = with_viewer(CausationHold(), port=8080)
        # Now all captures show in browser
    """
    bridge = HoldBridge(hold, port=port)
    bridge.start()
    return hold


# Convenience runner for testing with real models
def run_with_viewer(
    model_fn: Callable,
    port: int = 8080,
    steps: int = 10,
):
    """
    Run a model with the HOLD viewer.
    
    Args:
        model_fn: Function that takes (hold, step) and returns (obs, probs, value)
        port: Viewer port
        steps: Number of inference steps
        
    Usage:
        def my_model(hold, step):
            obs = env.reset() if step == 0 else env.step(last_action)
            probs = policy(obs)
            value = critic(obs)
            return obs, probs, value
            
        run_with_viewer(my_model, steps=20)
    """
    from .session import CausationHold
    
    hold = CausationHold()
    bridge = HoldBridge(hold, port=port)
    bridge.start()
    
    session = hold.begin_session("viewer_demo")
    
    try:
        for step in range(steps):
            obs, probs, value = model_fn(hold, step)
            
            # Convert to candidates
            candidates = [
                {'value': i, 'probability': float(p)}
                for i, p in enumerate(probs)
            ]
            
            choice, feedback = hold.capture(
                input_context={'observation': obs, 'step': step},
                candidates=candidates,
            )
            
            print(f"Step {step}: action={choice}, override={feedback.was_override if feedback else False}")
            
    finally:
        stats = hold.end_session()
        bridge.stop()
        print(f"\nSession stats: {stats.total_captures} captures, {stats.overrides} overrides")
