"""
Embodied runtime adapter.

This module owns local sensors/effectors and delegates decision making to
ActionPolicy. The shared state remains the workspace vector plus serializable
events, so memory and runtime orchestration can evolve independently.
"""
from __future__ import annotations

import torch
import numpy as np
import time
from pynput import mouse as pynput_mouse
from pynput import keyboard as pynput_keyboard

from .consensus import ConsensusSnapshot
from .events import WorkspaceEvent
from .memory import Hippocampus, InMemoryVectorMemoryStore
from .policy import (
    ActionPolicy,
    BasalGanglia,
    Cerebellum,
    CentralNervousSystem,
    ReflexArc,
    compose_action_params,
)

class MouseActuator:
    """
    鼠标执行器——对应"手"。

    workspace 解码出动作向量 → 移动鼠标 + 点击。
    动作空间：
        (dx, dy, click_left, click_right, scroll)
        dx, dy: 相对移动量（-1 到 1，乘以灵敏度）
        click_left/right: 0 或 1
        scroll: 滚动量
    """

    def __init__(self, sensitivity: int = 200, enabled: bool = True):
        self.sensitivity = sensitivity
        self.enabled = enabled
        self.controller = pynput_mouse.Controller() if enabled else None

    def execute(self, action: np.ndarray):
        """执行鼠标动作

        Args:
            action: (5,) = (dx, dy, click_l, click_r, scroll)
        """
        if not self.enabled or self.controller is None:
            return

        dx, dy, click_l, click_r, scroll = action

        # 移动
        if abs(dx) > 0.01 or abs(dy) > 0.01:
            self.controller.move(int(dx * self.sensitivity), int(dy * self.sensitivity))

        # 点击
        if click_l > 0.5:
            self.controller.click(pynput_mouse.Button.left)
            time.sleep(0.05)
        if click_r > 0.5:
            self.controller.click(pynput_mouse.Button.right)
            time.sleep(0.05)

        # 滚动
        if abs(scroll) > 0.1:
            self.controller.scroll(0, int(scroll * 5))

    def get_position(self) -> tuple[int, int]:
        if self.controller:
            return self.controller.position
        return (0, 0)


class KeyboardActuator:
    """
    键盘执行器——对应"手"+"口"（打字）。

    workspace 解码出 token → 按键输入。
    """

    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self.controller = pynput_keyboard.Controller() if enabled else None

    def type_text(self, text: str):
        """输入文本"""
        if not self.enabled or self.controller is None:
            return
        self.controller.type(text)

    def press_key(self, key: str):
        """按下单个键"""
        if not self.enabled or self.controller is None:
            return
        try:
            self.controller.press(key)
            self.controller.release(key)
        except Exception:
            pass


class AudioActuator:
    """
    音频执行器——对应"口"（说话）。

    workspace 解码出音频波形 → 扬声器播放。
    """

    def __init__(self, sample_rate: int = 16000, enabled: bool = True):
        self.sample_rate = sample_rate
        self.enabled = enabled
        import sounddevice as sd
        self._sd = sd

    def play(self, audio: np.ndarray, blocking: bool = False):
        """播放音频"""
        if not self.enabled:
            return
        audio = np.clip(audio, -1, 1).astype(np.float32)
        self._sd.play(audio, self.sample_rate)
        if blocking:
            self._sd.wait()

    def stop(self):
        if self.enabled:
            self._sd.stop()


class ScreenActuator:
    """
    屏幕执行器——对应"手"（绘制）。

    在屏幕上绘制 workspace 解码出的图像。
    用 OpenCV 显示一个窗口。
    """

    def __init__(self, enabled: bool = True, window_name: str = "JspaceAI Output"):
        self.enabled = enabled
        self.window_name = window_name
        import cv2
        self._cv2 = cv2

    def show_image(self, img: np.ndarray):
        """显示图像"""
        if not self.enabled:
            return
        # img: (H, W, 3) RGB 或 (H, W) 灰度
        if img.ndim == 3 and img.shape[2] == 3:
            img_bgr = self._cv2.cvtColor(img.astype(np.uint8), self._cv2.COLOR_RGB2BGR)
        else:
            img_bgr = img.astype(np.uint8)
        self._cv2.imshow(self.window_name, img_bgr)
        self._cv2.waitKey(1)

    def close(self):
        if self.enabled:
            self._cv2.destroyWindow(self.window_name)

# ============================================================
# 完整的具身 Agent
# ============================================================

class EmbodiedAgent:
    """
    完整的具身 Agent——感知-思考-行动闭环。

    结构：
        感知层（眼耳皮肤）→ FullSensoryStream
            ↓ 编码
        大脑皮层（思考）→ MultimodalJSpaceModel
            ↓ workspace w
        海马体（记忆）→ Hippocampus 存储和检索
            ↓
        ActionPolicy（价值 + 电机参数 + 门控）
            ↓
        执行器（手足口）→ MouseActuator + KeyboardActuator + AudioActuator

    循环：
        1. 感知：从五感获取输入
        2. 思考：workspace 更新
        3. 回忆：海马体检索相关记忆
        4. 决策：ActionPolicy 选动作并门控
        7. 行动：执行器执行
        8. 学习：更新策略与海马体
    """

    def __init__(self, model, device: str = 'cpu',
                 enable_mouse_output: bool = False,
                 enable_keyboard_output: bool = False,
                 enable_audio_output: bool = True,
                 enable_screen_output: bool = True,
                 enable_memory: bool = True,
                 risk_threshold: float = 0.3):
        self.model = model
        self.device = device
        self.config = model.config

        # 感知层
        from .desktop import FullSensoryStream
        from .platform import get_screen_size
        self.screen_w, self.screen_h = get_screen_size()
        self.senses = FullSensoryStream(
            use_camera=True, use_mic=True, use_desktop=True,
            img_size=(self.config.img_size, self.config.img_size),
        )

        # 执行器层
        self.mouse_actuator = MouseActuator(enabled=enable_mouse_output)
        self.keyboard_actuator = KeyboardActuator(enabled=enable_keyboard_output)
        self.audio_actuator = AudioActuator(enabled=enable_audio_output)
        self.screen_actuator = ScreenActuator(enabled=enable_screen_output)

        # 策略层
        self.policy = ActionPolicy(
            workspace_dim=self.config.workspace_dim,
            action_dim=5,
            n_actions=5,
            base_threshold=risk_threshold,
            device=device,
        )
        self.cerebellum = self.policy.motor_controller
        self.cns = self.policy.gate
        self.hippocampus = InMemoryVectorMemoryStore(
            workspace_dim=self.config.workspace_dim,
        ) if enable_memory else None
        self.basal_ganglia = self.policy.value_model

        # 内部状态
        self.state = model.init_state(1, torch.device(device))
        self.step_count = 0
        self.running = False
        self.risk_threshold = risk_threshold
        self.last_consensus: ConsensusSnapshot | None = None

        # 设置反射弧
        self._setup_reflexes()

    def _setup_reflexes(self):
        """设置基本反射"""
        # 反射1：大声音 → 抑制动作
        def loud_noise_condition(state):
            audio = state.get('audio')
            if audio and hasattr(audio, 'data'):
                return np.abs(audio.data).mean() > 0.5
            return False

        def loud_noise_response():
            self.cns.inhibit_score = 0.5

        self.policy.add_reflex(ReflexArc(
            trigger='loud_noise',
            condition=loud_noise_condition,
            action=loud_noise_response,
            priority=10,
        ))

    def perceive(self) -> dict:
        """感知：从五感获取输入"""
        return self.senses.get_latest()

    def think(self, sensory_data: dict) -> tuple[torch.Tensor, str]:
        """思考：更新 workspace

        Returns:
            w: workspace 状态
            modality: 本次输入的模态
        """
        # 按优先级选模态
        modality = None
        input_tensor = None

        if sensory_data['keyboard']:
            key_ids = [ord(ev.data[0]) if ev.data and len(ev.data) == 1 else 0
                       for ev in sensory_data['keyboard'][:8]]
            if key_ids:
                modality = 'keyboard'
                input_tensor = torch.tensor([key_ids], dtype=torch.long).to(self.device)

        elif sensory_data['mouse']:
            ev = sensory_data['mouse'][0]
            x, y = ev.data
            modality = 'mouse'
            input_tensor = torch.tensor([[
                x / self.screen_w, y / self.screen_h,
                1.0 if ev.modifiers.get('button') == 'left' else 0.0,
                1.0 if ev.modifiers.get('button') == 'right' else 0.0,
            ]], dtype=torch.float32).to(self.device)

        elif sensory_data['audio']:
            modality = 'audio'
            input_tensor = torch.tensor([sensory_data['audio'].data],
                                       dtype=torch.float32).to(self.device)

        elif sensory_data['screen']:
            img = sensory_data['screen'].data
            modality = 'screen'
            input_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0

        elif sensory_data['camera']:
            img = sensory_data['camera'].data
            modality = 'image'
            input_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0

        if modality is None or input_tensor is None:
            return self.state['w'], 'idle'

        # 编码 + forward
        with torch.no_grad():
            x = self.model.encode_modality(modality, input_tensor)
            if x.dim() == 1: x = x.unsqueeze(0)
            if x.dim() == 3: x = x[:, -1, :]
            if x.shape[0] != 1: x = x[-1:]
            self.state, alpha, _ = self.model.step(self.state, x)
            self.last_consensus = ConsensusSnapshot.from_workspace(
                self.state['w'], alpha, modality, self.model.expert_modality
            )

        return self.state['w'], modality

    def remember(self, w: torch.Tensor, context: dict):
        """存储 workspace event 到记忆层"""
        if self.hippocampus:
            event = WorkspaceEvent.from_tensor(
                w,
                modality=str(context.get('modality', 'unknown')),
                step=int(context.get('step', self.step_count)),
                consensus=context.get('consensus', self.last_consensus),
                context=context,
            )
            self.hippocampus.put(event)

    def recall_memories(self, w: torch.Tensor, top_k: int = 3) -> list:
        """从海马体回忆"""
        if self.hippocampus:
            return self.hippocampus.recall(w[0].cpu().numpy(), top_k)
        return []

    def decide_and_act(self, w: torch.Tensor, modality: str) -> dict:
        """决策和行动

        1. 基底神经节选动作
        2. 小脑计算动作参数
        3. 中枢神经门控
        4. 执行器执行
        """
        decision = self.policy.decide(w)
        action_params = decision.action_params
        execute = decision.should_execute

        # 4. 执行
        if execute:
            # 鼠标动作
            self.mouse_actuator.execute(action_params)
            # 音频输出（从 workspace 解码）
            if self.step_count % 10 == 0:
                with torch.no_grad():
                    audio_out = self.model.audio_decoder(w)[0].cpu().numpy()
                self.audio_actuator.play(audio_out)
            # 屏幕输出
            if self.screen_actuator.enabled:
                with torch.no_grad():
                    img_out = self.model.visual_decoder(w)[0].cpu().permute(1, 2, 0).numpy()
                    img_out = ((img_out + 1) / 2 * 255).clip(0, 255).astype(np.uint8)
                    self.screen_actuator.show_image(img_out)

            self.policy.record_action(action_params, modality)

        return decision.to_dict(executed=execute)

    def learn(self, w: torch.Tensor, action_idx: int, reward: float = 0.0):
        """学习——更新基底神经节和小脑"""
        self.policy.learn(w, action_idx, reward)

    def step_once(self) -> dict:
        """执行一步完整的感知-思考-行动循环"""
        # 1. 感知
        sensory_data = self.perceive()

        # 2. 检查反射
        reflex_action = self.cns.check_reflexes(sensory_data)
        if reflex_action:
            reflex_action()

        # 3. 思考
        w, modality = self.think(sensory_data)

        # 4. 回忆
        self.recall_memories(w)

        # 5. 决策和行动
        action_info = self.decide_and_act(w, modality)

        # 6. 学习（自监督：预测误差作为 reward）
        reward = -action_info['risk']  # 简化：风险越低 reward 越高
        self.learn(w, action_info['action_idx'], reward)

        # 7. 记忆存储
        self.remember(w, {
            'modality': modality,
            'action': action_info,
            'step': self.step_count,
        })

        self.step_count += 1

        return {
            'step': self.step_count,
            'modality': modality,
            'w_norm': w.norm().item(),
            'action': action_info,
            'consensus': self.last_consensus.to_dict() if self.last_consensus else None,
            'memories_count': self.hippocampus.size() if self.hippocampus else 0,
        }

    def run(self, n_steps: int = 100, interval: float = 0.2,
            on_step: callable = None):
        """运行感知-思考-行动循环"""
        self.running = True
        self.senses.start()
        print(f"具身 Agent 启动，{n_steps} 步")
        print(f"执行器: mouse={self.mouse_actuator.enabled}, "
              f"keyboard={self.keyboard_actuator.enabled}, "
              f"audio={self.audio_actuator.enabled}, "
              f"screen={self.screen_actuator.enabled}")
        print("=" * 60)

        step_log = []
        try:
            for _ in range(n_steps):
                if not self.running:
                    break
                info = self.step_once()
                step_log.append(info)

                if on_step:
                    on_step(info)
                elif info['step'] % 5 == 0:
                    print(f"  step {info['step']:3d} | mod {info['modality']:8s} | "
                          f"||w|| {info['w_norm']:.3f} | "
                          f"action {info['action']['action_name']} | "
                          f"executed {info['action']['executed']}")

                time.sleep(interval)
        except KeyboardInterrupt:
            print("\n用户中断")
        finally:
            self.running = False
            self.senses.stop()
            self.audio_actuator.stop()
            if self.screen_actuator.enabled:
                self.screen_actuator.close()

        return step_log