"""
Tiered Brain Architecture for OpenCastor.

Three layers, fastest first:

  Layer 0 — Reactive (rule-based, <1ms)
    Hardcoded safety: obstacle too close → stop, blank frame → wait,
    e-stop → halt. No LLM needed.

  Layer 1 — Fast Brain (Gemini Flash / Ollama, ~1-2s)
    Primary perception-action loop. Processes camera frames,
    produces JSON actions. Handles routine navigation.

  Layer 2 — Planner (Claude / Opus, ~10-15s)
    Complex reasoning, scene understanding, conversation,
    multi-step planning. Called periodically or on escalation.

The control loop runs Layer 0 every tick, Layer 1 every tick (async),
and Layer 2 every N ticks or when Layer 1 signals uncertainty.
"""

import logging
import time

from .providers.base import Thought
from .providers.task_router import TaskCategory, TaskRouter

logger = logging.getLogger("OpenCastor.TieredBrain")

# Task categories that should ALWAYS skip the planner (cheap/fast tier only)
_FAST_ONLY_CATEGORIES: frozenset[TaskCategory] = frozenset({TaskCategory.SENSOR_POLL})

# Task categories that should PREFER the planner when available
_PREFER_PLANNER_CATEGORIES: frozenset[TaskCategory] = frozenset(
    {
        TaskCategory.REASONING,
        TaskCategory.CODE,
        TaskCategory.SAFETY,
        TaskCategory.VISION,
        TaskCategory.SEARCH,
    }
)


class ReactiveLayer:
    """Layer 0: Rule-based reactive safety controller.

    Combines hardcoded rules (<1ms) with optional Hailo-8 NPU
    object detection (~20ms) for obstacle avoidance without API calls.
    Returns an action if triggered, None to pass to next layer.
    """

    def __init__(self, config: dict):
        reactive = config.get("reactive", {})
        self.min_obstacle_m = reactive.get("min_obstacle_m", 0.3)
        self.blank_threshold = reactive.get("blank_threshold", 100)
        self.hailo_enabled = reactive.get("hailo_vision", False)
        # Distance thresholds for Hailo-8 NPU detections.
        # hailo_stop_distance_m: e-stop when nearest obstacle is closer than this.
        # hailo_warn_distance_m: slow down / avoid when closer than this.
        self.hailo_stop_distance_m = reactive.get("hailo_stop_distance_m", 0.5)
        self.hailo_warn_distance_m = reactive.get("hailo_warn_distance_m", 1.0)
        self.hailo_calibration = reactive.get("hailo_calibration", 0.25)
        # If camera_required=False, blank/missing frames are NOT a blocking condition.
        # The brain will run text-only (messaging, sensor data) without a live frame.
        self.camera_required = config.get("camera", {}).get("camera_required", True)
        self._hailo = None
        self.last_detections = []  # Expose for telemetry/logging

        if self.hailo_enabled:
            try:
                from .hailo_vision import HailoVision

                model = config.get("reactive", {}).get(
                    "hailo_model", "/usr/share/hailo-models/yolov8s_h8.hef"
                )
                confidence = config.get("reactive", {}).get("hailo_confidence", 0.4)
                self._hailo = HailoVision(model_path=model, confidence=confidence)
                if not self._hailo.available:
                    self._hailo = None
            except Exception as e:
                logger.debug(f"Hailo vision not available: {e}")

    def evaluate(self, frame_bytes: bytes, sensor_data: dict | None = None) -> dict | None:
        """Check reactive safety rules. Returns action dict or None."""
        # Rule 1: Blank/missing frame → wait (skipped if camera_required=False)
        if not frame_bytes or len(frame_bytes) < self.blank_threshold:
            if self.camera_required:
                return {"type": "wait", "duration_ms": 500, "reason": "no_camera_data"}
            # camera_required=False: pass through to fast brain (text/sensor-only mode)
            return None

        # Rule 2: All-black frame (camera blocked/failed) — skipped if camera_required=False
        if frame_bytes == b"\x00" * len(frame_bytes):
            if self.camera_required:
                return {"type": "wait", "duration_ms": 500, "reason": "blank_frame"}
            return None

        # Rule 3: Depth-based obstacle proximity
        if sensor_data:
            distance = sensor_data.get("front_distance_m")
            if distance is not None and distance < self.min_obstacle_m:
                logger.warning(f"Reactive: obstacle at {distance:.2f}m — stopping!")
                return {"type": "stop", "reason": f"obstacle_{distance:.2f}m"}

        # Rule 4: Battery critical
        if sensor_data and sensor_data.get("battery_critical"):
            return {"type": "stop", "reason": "battery_critical"}

        # Rule 5: Hailo-8 NPU object detection (~20ms)
        if self._hailo is not None:
            try:
                import cv2
                import numpy as np

                # Decode JPEG to frame
                arr = np.frombuffer(frame_bytes, dtype=np.uint8)
                frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
                if frame is not None:
                    result = self._hailo.detect_obstacles(frame)
                    self.last_detections = result.get("all_detections", [])

                    nearest = result.get("nearest_obstacle")
                    if nearest:
                        dist_m = nearest.estimate_distance_m(self.hailo_calibration)
                        if dist_m <= self.hailo_stop_distance_m:
                            logger.warning(
                                "Reactive: %s at ~%.2fm — e-stop!",
                                nearest.class_name,
                                dist_m,
                            )
                            return {
                                "type": "stop",
                                "reason": f"hailo_{nearest.class_name}_{dist_m:.2f}m",
                            }
                        if dist_m <= self.hailo_warn_distance_m:
                            logger.info(
                                "Reactive: %s at ~%.2fm — slowing",
                                nearest.class_name,
                                dist_m,
                            )
                            return {
                                "type": "move",
                                "linear": 0.0,
                                "angular": 0.3,
                                "reason": f"hailo_warn_{nearest.class_name}_{dist_m:.2f}m",
                            }

                    if not result["clear_path"] and result["obstacles"]:
                        # Obstacles in center path but beyond warn distance — nudge
                        names = [d.class_name for d in result["obstacles"][:3]]
                        return {
                            "type": "move",
                            "linear": 0.0,
                            "angular": 0.3,  # Turn to avoid
                            "reason": f"hailo_avoid_{','.join(names)}",
                        }
            except Exception as e:
                logger.debug(f"Hailo detection error: {e}")

        # No reactive trigger — pass to next layer
        return None

    def close(self):
        """Release Hailo resources."""
        if self._hailo:
            self._hailo.close()


class TieredBrain:
    """Orchestrates the three brain layers.

    The fast brain runs every tick. The planner runs every
    `planner_interval` ticks or when the fast brain signals
    uncertainty (action confidence < threshold).
    """

    def __init__(
        self,
        fast_provider,
        planner_provider=None,
        config: dict | None = None,
        task_router: TaskRouter | None = None,
    ):
        config = config or {}
        self.fast = fast_provider
        self.planner = planner_provider
        self.reactive = ReactiveLayer(config)

        # Task-aware routing (issue #612): optional TaskRouter for category-based dispatch.
        # If a task_routing block exists in RCAN config, honour it; else use defaults.
        routing_table = config.get("task_routing", {})
        self.task_router: TaskRouter = task_router or TaskRouter(
            routing_table=routing_table or None
        )

        # Planner runs every N ticks (0 = never auto-run)
        self.planner_interval = config.get("tiered_brain", {}).get("planner_interval", 10)
        self.uncertainty_threshold = config.get("tiered_brain", {}).get(
            "uncertainty_threshold", 0.3
        )
        self.tick_count = 0
        self.last_plan = None
        self.last_plan_time = 0

        # Layer 3: Agent Swarm (optional — enabled via agents.enabled: true in RCAN config)
        self.orchestrator = None
        if config.get("agents", {}).get("enabled", False):
            try:
                from .agents.orchestrator import OrchestratorAgent
                from .agents.shared_state import SharedState

                agent_cfg = config.get("agents", {})
                self.orchestrator = OrchestratorAgent(
                    config=agent_cfg,
                    shared_state=SharedState(),
                )
                logger.info("Layer 3 (Agent Swarm) enabled")
            except Exception as exc:
                logger.debug("Layer 3 not available: %s", exc)

        # Embedding Interpreter (optional — config-gated, best-effort)
        self.interpreter = None
        if config.get("interpreter", {}).get("enabled", False):
            try:
                from .embedding_interpreter import EmbeddingInterpreter

                self.interpreter = EmbeddingInterpreter(config.get("interpreter", {}))
                logger.info(
                    "EmbeddingInterpreter enabled (backend=%s)",
                    config.get("interpreter", {}).get("backend", "auto"),
                )
            except Exception as exc:
                logger.warning("EmbeddingInterpreter not available: %s", exc)

        # Stats
        self.stats = {
            "reactive_count": 0,
            "fast_count": 0,
            "planner_count": 0,
            "swarm_count": 0,
            "total_ticks": 0,
            "interpreter_pre_count": 0,
            "interpreter_escalations": 0,
        }

    def think(
        self,
        image_bytes: bytes,
        instruction: str,
        sensor_data: dict | None = None,
        task_category: str | None = None,
    ) -> Thought:
        """Run the tiered brain pipeline.

        Args:
            image_bytes:    Raw camera frame bytes.
            instruction:    Natural-language task description.
            sensor_data:    Optional sensor snapshot dict.
            task_category:  Optional :class:`~castor.providers.task_router.TaskCategory`
                            value (string) for task-aware routing (issue #612).

                            - ``"sensor_poll"`` — never escalate to planner; cheap/fast only.
                            - ``"reasoning"``, ``"code"``, ``"safety"``, ``"vision"``,
                              ``"search"`` — prefer planner when available.
                            - ``"navigation"`` or ``None`` — default interval-based behaviour.

        Returns:
            A :class:`~castor.providers.base.Thought` with an action and metadata.
        """
        self.tick_count += 1
        self.stats["total_ticks"] += 1

        # ── Task-category routing (issue #612) ──────────────────────────────
        _resolved_category: TaskCategory | None = None
        if task_category is not None:
            try:
                _resolved_category = TaskCategory(task_category)
                logger.debug("TieredBrain: task_category=%s", _resolved_category.value)
            except ValueError:
                logger.warning("Unknown task_category %r — ignoring for routing", task_category)

        # Layer 0: Reactive (instant)
        reactive_action = self.reactive.evaluate(image_bytes, sensor_data)
        if reactive_action:
            self.stats["reactive_count"] += 1
            logger.debug(
                f"Reactive: {reactive_action['type']} ({reactive_action.get('reason', '')})"
            )
            return Thought(f"Reactive: {reactive_action.get('reason', '')}", reactive_action)

        # Layer 2 escalation flag — may be set early by interpreter
        should_plan = False

        # Interpreter pre-think (non-blocking — catches all exceptions)
        scene_ctx = None
        if self.interpreter:
            try:
                scene_ctx = self.interpreter.pre_think(image_bytes, instruction, sensor_data)
                self.stats["interpreter_pre_count"] += 1
                if scene_ctx.should_escalate:
                    should_plan = True
                    self.stats["interpreter_escalations"] += 1
                    logger.info(
                        "Interpreter: goal_similarity=%.2f → forcing L2 escalation",
                        scene_ctx.goal_similarity,
                    )
            except Exception as exc:
                logger.debug("Interpreter pre_think (non-fatal): %s", exc)

        # Layer 1: Fast brain
        t0 = time.time()
        thought = self.fast.think(image_bytes, instruction)
        fast_ms = (time.time() - t0) * 1000
        thought.layer = "fast"
        self.stats["fast_count"] += 1

        if thought.action:
            logger.info(f"Fast brain ({fast_ms:.0f}ms): {thought.action.get('type', '?')}")

        # Layer 2: Planner (periodic or on escalation)
        # Task-category overrides: SENSOR_POLL skips planner; high-complexity categories force it.
        if _resolved_category in _FAST_ONLY_CATEGORIES:
            # Cheap tasks: NEVER escalate to planner — hard override.
            # Explicitly clear should_plan even if interpreter pre-think set should_escalate=True,
            # so sensor_poll ticks stay cheap regardless of prior escalation signals.
            should_plan = False
            logger.debug(
                "TieredBrain: task_category=%s → fast-only, planner suppressed",
                _resolved_category.value if _resolved_category else None,
            )
        elif _resolved_category in _PREFER_PLANNER_CATEGORIES and self.planner:
            # High-complexity tasks: force planner when available
            should_plan = True
            logger.info("TieredBrain: task_category=%s → forcing planner", _resolved_category.value)
        else:
            # Default interval-based logic (covers NAVIGATION, None, unknown)
            if self.planner and self.planner_interval > 0:
                if self.tick_count % self.planner_interval == 0:
                    should_plan = True
                    logger.info("Planner: periodic check (tick %d)", self.tick_count)

            # Also escalate if fast brain produced no action
            if self.planner and not thought.action:
                should_plan = True
                logger.info("Planner: escalation (fast brain produced no action)")

        if should_plan and self.planner:
            try:
                # Inject dynamic sensor state into the USER message (not the system prompt).
                # This keeps the system prompt prefix stable across ticks so cache hits occur.
                # Per Claude Code's cache-first lesson: static content in system, dynamic in user.
                from castor.prompt_cache import build_sensor_reminder

                sensor_reminder = build_sensor_reminder(sensor_data or {})
                plan_instruction = (f"{sensor_reminder}\n\n" if sensor_reminder else "") + (
                    f"You are the strategic planner for a robot. "
                    f"The fast brain's last response: {thought.raw_text[:200]}\n\n"
                    f"Current task: {instruction}\n\n"
                    f"Provide a high-level plan or corrected action as JSON."
                )

                # Inject RAG context from embedding interpreter
                if scene_ctx and self.interpreter:
                    try:
                        rag = self.interpreter.format_rag_context(scene_ctx)
                        if rag:
                            plan_instruction = rag + "\n\n" + plan_instruction
                    except Exception:
                        pass

                t0 = time.time()
                plan_thought = self.planner.think(image_bytes, plan_instruction)
                plan_ms = (time.time() - t0) * 1000
                plan_thought.layer = "planner"
                plan_thought.escalated = True
                self.stats["planner_count"] += 1
                if plan_thought.action:
                    self.last_plan = plan_thought.action
                    self.last_plan_time = time.time()
                    logger.info(
                        f"Planner ({plan_ms:.0f}ms): {plan_thought.action.get('type', '?')}"
                    )
                    # Planner overrides fast brain when it has a plan
                    # Interpreter post-think: store episode for planner path
                    if self.interpreter and scene_ctx:
                        try:
                            self.interpreter.post_think(scene_ctx, plan_thought)
                        except Exception as exc:
                            logger.debug("Interpreter post_think (non-fatal): %s", exc)
                    return plan_thought
            except Exception as e:
                logger.warning(f"Planner error (non-fatal): {e}")

        # Layer 3: Agent Swarm (async orchestration — only if enabled)
        if self.orchestrator is not None:
            try:
                swarm_action = self.orchestrator.sync_think(sensor_data or {})
                if swarm_action.get("type") not in (None, "idle"):
                    self.stats["swarm_count"] += 1
                    logger.debug("Layer 3 swarm action: %s", swarm_action.get("type"))
                    swarm_thought = Thought(
                        f"Swarm: {swarm_action.get('type', '?')}",
                        swarm_action,
                    )
                    # Interpreter post-think: store episode for swarm path
                    if self.interpreter and scene_ctx:
                        try:
                            self.interpreter.post_think(scene_ctx, swarm_thought)
                        except Exception as exc:
                            logger.debug("Interpreter post_think (non-fatal): %s", exc)
                    return swarm_thought
            except Exception as exc:
                logger.debug("Layer 3 error (non-fatal): %s", exc)

        # Interpreter post-think: store episode
        if self.interpreter and scene_ctx:
            try:
                self.interpreter.post_think(scene_ctx, thought)
            except Exception as exc:
                logger.debug("Interpreter post_think (non-fatal): %s", exc)

        return thought

    def get_stats(self) -> dict:
        """Return brain layer usage stats."""
        total = max(self.stats["total_ticks"], 1)
        stats = {
            **self.stats,
            "reactive_pct": round(self.stats["reactive_count"] / total * 100, 1),
            "fast_pct": round(self.stats["fast_count"] / total * 100, 1),
            "planner_pct": round(self.stats["planner_count"] / total * 100, 1),
            "swarm_pct": round(self.stats["swarm_count"] / total * 100, 1),
        }
        # Include prompt cache stats from planner if available
        if self.planner and hasattr(self.planner, "cache_stats"):
            stats["cache"] = self.planner.cache_stats
        if self.interpreter and hasattr(self.interpreter, "status"):
            stats["interpreter"] = self.interpreter.status()
        return stats
