"""Base classes for robot behaviors."""

from __future__ import annotations

import os
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, TYPE_CHECKING

from .config import VoiceConfig, LLMConfig
from .robot import RobotAdapter
from .state import RobotStep, SharedState
from .voice import rule_parse

if TYPE_CHECKING:
    from .scene import SceneState, SceneObject

# Try to import dspy (optional dependency)
try:
    import dspy
    DSPY_AVAILABLE = True
except ImportError:
    dspy = None  # type: ignore
    DSPY_AVAILABLE = False


def _log(msg: str) -> None:
    """Print a timestamped log message."""
    timestamp = time.strftime("%H:%M:%S")
    print(f"[behaviors {timestamp}] {msg}", flush=True)


@dataclass
class ActionContext:
    """Context for action execution."""

    pose: Optional[List[float]]  # [x, y, z, roll, pitch, yaw]
    target: Optional["SceneObject"]  # Target object
    scene: "SceneState"  # Scene state with all objects
    config: VoiceConfig
    home_pose: Dict[str, float]
    shared_state: SharedState


@dataclass
class ActionInfo:
    """Metadata for an action."""

    name: str
    aliases: List[str] = field(default_factory=list)
    requires_object: bool = False
    requires_pose: bool = False
    description: str = ""


# Default actions for pick-and-place robots
DEFAULT_ACTIONS: Dict[str, ActionInfo] = {
    "subir": ActionInfo(
        name="subir",
        aliases=["sube", "arriba"],
        requires_pose=True,
        description="Subir el robot",
    ),
    "bajar": ActionInfo(
        name="bajar",
        aliases=["baja", "abajo"],
        requires_pose=True,
        description="Bajar el robot",
    ),
    "ir": ActionInfo(
        name="ir",
        aliases=["ve", "mover", "muevete", "acercar"],
        requires_object=True,
        description="Ir hacia un objeto",
    ),
    "tomar": ActionInfo(
        name="tomar",
        aliases=["toma", "agarra", "agarrar", "coger", "chupar", "succionar"],
        requires_object=True,
        description="Tomar un objeto",
    ),
    "soltar": ActionInfo(
        name="soltar",
        aliases=["deja", "dejar"],
        requires_object=True,
        description="Soltar el objeto",
    ),
    "reiniciar": ActionInfo(
        name="reiniciar",
        aliases=["reinicia", "reset"],
        requires_pose=False,
        requires_object=False,
        description="Reiniciar a posicion inicial",
    ),
}


class RobotBehavior(ABC):
    """
    Robot behavior with configurable robot adapter.

    Uses a RobotAdapter to translate generic commands (grab, release, move)
    to robot-specific commands (vacuum_on, gripper_close, etc.).
    """

    # Class-level action definitions
    ACTIONS: ClassVar[Dict[str, ActionInfo]] = DEFAULT_ACTIONS

    # Supported language
    LANGUAGE: ClassVar[str] = "es"

    # DSPy signature class (override in subclasses)
    CommandSignature: ClassVar[type] = None  # type: ignore

    def __init__(
        self,
        config: VoiceConfig,
        robot_adapter: RobotAdapter,
        llm_config: Optional[LLMConfig] = None,
    ) -> None:
        self.config = config
        self.robot_adapter = robot_adapter
        self._alias_map: Dict[str, str] = {}
        self._predictor: Optional[Any] = None
        self._build_alias_map()
        if llm_config and llm_config.provider != "rules":
            self._init_dspy(llm_config)

    def _init_dspy(self, llm_config: LLMConfig) -> None:
        """Initialize DSPy predictor for this behavior."""
        _log(f"Initializing DSPy with provider={llm_config.provider}, model={llm_config.model}")
        if not DSPY_AVAILABLE:
            _log("DSPy not available (import failed), falling back to rules")
            return
        if self.CommandSignature is None:
            _log("No CommandSignature defined, falling back to rules")
            return
        try:
            lm = self._create_lm(llm_config)
            if lm:
                dspy.configure(lm=lm)
                self._predictor = dspy.Predict(self.CommandSignature)
                _log(f"DSPy initialized with {llm_config.provider}/{llm_config.model}")
            else:
                _log(f"Failed to create LM for provider={llm_config.provider}")
        except Exception as e:
            import traceback
            _log(f"Failed to initialize DSPy: {e}")
            _log(f"Traceback: {traceback.format_exc()}")

    def _create_lm(self, config: LLMConfig) -> Optional[Any]:
        """Create DSPy language model."""
        if not DSPY_AVAILABLE:
            return None
        if config.provider == "gemini":
            api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
            return dspy.LM(
                f"gemini/{config.model}",
                api_key=api_key,
                temperature=config.temperature,
            )
        if config.provider == "ollama":
            return dspy.LM(
                f"ollama_chat/{config.model}",
                api_base=config.api_base,
                api_key="",
            )
        if config.provider == "openai":
            return dspy.LM(
                f"openai/{config.model}",
                temperature=config.temperature,
            )
        return None

    def parse_command(self, transcript: str) -> Dict[str, str]:
        """Parse voice command using DSPy or fallback to rules."""
        if self._predictor:
            _log(f"Using DSPy to parse: '{transcript}'")
            try:
                result = self._predictor(comando=transcript)
                _log(f"DSPy result: accion={result.accion}, objeto={result.objeto}, color={result.color}, tamano={result.tamano}")
                return {
                    "resultado": "ok" if result.accion != "error" else "error",
                    "accion": result.accion,
                    "objeto": result.objeto,
                    "color": result.color,
                    "tamano": result.tamano,
                }
            except Exception as e:
                import traceback
                _log(f"DSPy parsing failed: {e}, falling back to rules")
                _log(f"Traceback: {traceback.format_exc()}")
        else:
            _log(f"No DSPy predictor, using rules to parse: '{transcript}'")
        return self.rule_parse(transcript)

    def rule_parse(self, transcript: str) -> Dict[str, str]:
        """Rule-based parsing fallback using behavior's vocabulary."""
        return rule_parse(transcript, self.get_all_keywords())

    def _build_alias_map(self) -> None:
        """Build reverse mapping from aliases to action names."""
        self._alias_map = {}
        for action_name, info in self.ACTIONS.items():
            self._alias_map[action_name] = action_name
            for alias in info.aliases:
                self._alias_map[alias] = action_name

    def resolve_action(self, name: str) -> Optional[str]:
        """Resolve action name or alias to canonical name."""
        return self._alias_map.get(name)

    def get_action_info(self, name: str) -> Optional[ActionInfo]:
        """Get action info by name or alias."""
        canonical = self.resolve_action(name)
        return self.ACTIONS.get(canonical) if canonical else None

    def get_all_keywords(self) -> Dict[str, List[str]]:
        """Get all action names and aliases for parser."""
        return {
            name: [name] + info.aliases
            for name, info in self.ACTIONS.items()
        }

    def execute(self, action_name: str, context: ActionContext) -> bool:
        """Execute action by name using explicit handler mapping."""
        canonical = self.resolve_action(action_name)
        if not canonical:
            return False

        info = self.ACTIONS.get(canonical)
        if not info:
            return False

        # Check preconditions
        if info.requires_pose and not context.pose:
            return False
        if info.requires_object and not context.target:
            return False

        handler = self.action_handlers().get(canonical)
        if not handler:
            return False

        return handler(context)

    @abstractmethod
    def action_handlers(self) -> Dict[str, Callable[[ActionContext], bool]]:
        """Return mapping of action names to handler functions."""
        pass

    def _within_bounds(
        self,
        point_mm: List[float],
        workspace_min: Tuple[Optional[float], Optional[float], Optional[float]],
        workspace_max: Tuple[Optional[float], Optional[float], Optional[float]],
    ) -> bool:
        """Check if position is within workspace bounds."""
        x, y, z = point_mm[:3]
        min_x, min_y, min_z = workspace_min
        max_x, max_y, max_z = workspace_max
        if min_x is not None and x < min_x:
            return False
        if max_x is not None and x > max_x:
            return False
        if min_y is not None and y < min_y:
            return False
        if max_y is not None and y > max_y:
            return False
        if min_z is not None and z < min_z:
            return False
        if max_z is not None and z > max_z:
            return False
        return True

    def _queue_steps(self, ctx: ActionContext, steps: List[RobotStep]) -> None:
        """Queue multiple robot steps."""
        for step in steps:
            ctx.shared_state.append_queue(step)

    def _make_pose_payload(self, x: float, y: float, z: float) -> Dict[str, float]:
        """Create pose payload with default orientation."""
        return {
            "x": x,
            "y": y,
            "z": z,
            "roll": self.config.default_roll,
            "pitch": self.config.default_pitch,
            "yaw": self.config.default_yaw,
        }

    def _queue_move(self, ctx: ActionContext, x: float, y: float, z: float) -> bool:
        """Queue a move command if within bounds."""
        if not self._within_bounds([x, y, z], self.config.workspace_min, self.config.workspace_max):
            return False
        payload = self._make_pose_payload(x, y, z)
        self._queue_steps(ctx, self.robot_adapter.move(payload))
        return True

    def _queue_approach_sequence(self, ctx: ActionContext, pos: List[float]) -> None:
        """Queue approach pose then target pose."""
        approach_z = pos[2] + self.config.tcp_offset_mm + self.config.approach_offset_mm
        target_z = pos[2] + self.config.tcp_offset_mm

        # Approach
        self._queue_move(ctx, pos[0], pos[1], approach_z)
        # Target
        self._queue_move(ctx, pos[0], pos[1], target_z)

    # =========================================================
    # Action implementations using robot adapter
    # =========================================================

    def action_subir(self, ctx: ActionContext) -> bool:
        """Move up by step_mm."""
        target_z = ctx.pose[2] + self.config.step_mm
        return self._queue_move(ctx, ctx.pose[0], ctx.pose[1], target_z)

    def action_bajar(self, ctx: ActionContext) -> bool:
        """Move down by step_mm."""
        target_z = ctx.pose[2] - self.config.step_mm
        return self._queue_move(ctx, ctx.pose[0], ctx.pose[1], target_z)

    def action_ir(self, ctx: ActionContext) -> bool:
        """Move to object position (approach + target)."""
        pos = ctx.target.position_mm
        self._queue_approach_sequence(ctx, pos)
        return True

    def action_tomar(self, ctx: ActionContext) -> bool:
        """Pick object using robot adapter."""
        pos = ctx.target.position_mm
        self._queue_approach_sequence(ctx, pos)
        self._queue_steps(ctx, self.robot_adapter.grab())
        return True

    def action_soltar(self, ctx: ActionContext) -> bool:
        """Release object at target using robot adapter."""
        pos = ctx.target.position_mm
        self._queue_approach_sequence(ctx, pos)
        self._queue_steps(ctx, self.robot_adapter.release())
        return True

    def action_reiniciar(self, ctx: ActionContext) -> bool:
        """Reset: release tool, move home, clear objects."""
        self._queue_steps(ctx, self.robot_adapter.reset_tool())
        self._queue_steps(ctx, self.robot_adapter.move(ctx.home_pose))
        ctx.scene.clear_detected()
        return True
