"""Mission orchestration helpers that combine robot controls with MQTT feedback.

This module expands on the base Husky robot example by exposing a higher level
command that executes a multi-step pick-and-place routine.  The routine chains
existing motion templates from :mod:`omnilink.examples.robot.use_case` and wraps
them in the operator confirmation workflow demonstrated in
:mod:`omnilink.examples.messaging_demo`.

The mission handler relies on the MQTT bridge messenger so that it can emit
structured progress updates and pause execution until an operator approves
high-risk manoeuvres.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Sequence, Tuple

from omnilink import AgentFeedback, AgentQuestion, OmniLinkEngine, OmniLinkMQTTBridge

from .api import stop
from .use_case import build_use_case

ApprovalReply = Dict[str, object]


@dataclass(frozen=True)
class MissionStep:
    """Description of a single motion step executed during the mission."""

    description: str
    command: str
    requires_confirmation: bool = False
    metadata: Optional[Dict[str, object]] = None


DEFAULT_STEPS: Tuple[MissionStep, ...] = (
    MissionStep(
        "Approach staging area",
        "move_forward_at_0.4_m/s_for_3_seconds",
    ),
    MissionStep(
        "Rotate towards pallet",
        "turn_left_at_1.0_rad/s_for_1.2_seconds",
    ),
    MissionStep(
        "Creep into pickup zone",
        "move_forward_at_0.25_m/s_for_2.5_seconds",
        requires_confirmation=True,
        metadata={"risk": "entering shared workspace"},
    ),
    MissionStep(
        "Secure payload",
        "stop",
    ),
    MissionStep(
        "Retreat with payload",
        "move_backward_at_0.35_m/s_for_2.0_seconds",
    ),
    MissionStep(
        "Rotate toward drop point",
        "turn_right_at_1.0_rad/s_for_1.4_seconds",
    ),
    MissionStep(
        "Advance to placement area",
        "move_forward_at_0.3_m/s_for_1.8_seconds",
        requires_confirmation=True,
        metadata={"risk": "placing near operator"},
    ),
    MissionStep(
        "Release payload",
        "stop",
    ),
)

MISSION_TEMPLATE = "run_pick_and_place_mission"


def _normalise_reply(reply: ApprovalReply) -> str:
    """Return a lower-cased textual answer extracted from ``reply``."""

    for key in ("answer", "choice", "response"):
        raw = reply.get(key)
        if isinstance(raw, str) and raw.strip():
            return raw.strip().lower()
    # Some integrations nest the answer under ``data``
    data = reply.get("data")
    if isinstance(data, dict):
        for key in ("answer", "choice", "response"):
            raw = data.get(key)
            if isinstance(raw, str) and raw.strip():
                return raw.strip().lower()
    return ""


def _safe_stop() -> None:
    try:
        stop()
    except Exception:
        # Failing to stop should not crash the mission handler; the caller will
        # already receive structured error feedback.
        pass


def _require_messenger(evt: Dict[str, object]) -> object:
    messenger = evt.get("messenger")
    if messenger is None:
        raise RuntimeError(
            "robot mission demo requires a messenger-aware bridge such as OmniLinkMQTTBridge"
        )
    return messenger


def build_mission_engine(
    *,
    steps: Sequence[MissionStep] = DEFAULT_STEPS,
    approval_timeout: float = 45.0,
) -> Tuple[OmniLinkEngine, Optional[Callable[[], str]]]:
    """Return an engine configured with the base robot use case plus mission flow."""

    config = build_use_case()
    templates = list(config.templates)
    if MISSION_TEMPLATE not in templates:
        templates.append(MISSION_TEMPLATE)

    engine = OmniLinkEngine(templates, types=config.types)
    total_steps = len(steps)
    context_provider = config.context_provider

    def _mission_handler(evt: Dict[str, object]) -> Dict[str, object]:
        messenger = _require_messenger(evt)
        messenger.acknowledge(
            "accepted",
            message="Pick-and-place mission accepted",
            data={"steps": total_steps},
        )
        completed: List[Dict[str, object]] = []

        for index, step in enumerate(steps, start=1):
            progress = (index - 1) / max(total_steps, 1)
            step_payload = {
                "step": index,
                "of": total_steps,
                "command": step.command,
            }
            if step.metadata:
                step_payload["metadata"] = dict(step.metadata)
            messenger.send_feedback(
                AgentFeedback(
                    message=f"Starting step {index}/{total_steps}: {step.description}",
                    progress=progress,
                    data=step_payload,
                )
            )

            if step.requires_confirmation:
                messenger.acknowledge(
                    "pending",
                    message=f"Awaiting approval for step {index}: {step.description}",
                    data=step_payload,
                )
                pending = messenger.ask_question(
                    AgentQuestion(
                        prompt=f"Approve high-risk move '{step.description}'?",
                        choices=["yes", "no"],
                        data=step_payload,
                    )
                )
                try:
                    reply = pending.wait(timeout=approval_timeout, cancel_on_timeout=True)
                except TimeoutError:
                    messenger.send_feedback(
                        AgentFeedback(
                            message=f"Timed out waiting for approval of step {index}",
                            kind="error",
                            ok=False,
                            data=step_payload,
                        )
                    )
                    _safe_stop()
                    return {
                        "ack": False,
                        "aborted_at": index,
                        "reason": "timeout",
                    }

                answer = _normalise_reply(reply)
                if answer not in {"yes", "y", "approved", "approve"}:
                    messenger.send_feedback(
                        AgentFeedback(
                            message=f"Operator declined step {index}",
                            kind="error",
                            ok=False,
                            data={**step_payload, "reply": reply},
                        )
                    )
                    _safe_stop()
                    return {
                        "ack": False,
                        "aborted_at": index,
                        "reason": "declined",
                        "reply": reply,
                    }
                messenger.acknowledge(
                    "approved",
                    message=f"Operator approved step {index}",
                    data=step_payload,
                )

            motion_result = engine.handle(
                step.command,
                meta={"mission_step": index, "mission_command": evt.get("command")},
                messenger=messenger,
            )
            motion_payload = motion_result.get("result", {})
            if not isinstance(motion_payload, dict) or not motion_payload.get("ack"):
                messenger.send_feedback(
                    AgentFeedback(
                        message=f"Motion command failed during step {index}",
                        kind="error",
                        ok=False,
                        data={**step_payload, "result": motion_result},
                    )
                )
                return {
                    "ack": False,
                    "aborted_at": index,
                    "reason": "motion_failed",
                    "result": motion_result,
                }

            completed.append({**step_payload, "result": motion_payload})
            if context_provider is not None:
                messenger.send_feedback(
                    AgentFeedback(
                        message="Pose update",
                        data={
                            "step": index,
                            "pose_payload": context_provider(),
                        },
                    )
                )
            messenger.send_feedback(
                AgentFeedback(
                    message=f"Completed step {index}/{total_steps}: {step.description}",
                    progress=index / max(total_steps, 1),
                    data={**step_payload, "result": motion_payload},
                )
            )

        messenger.acknowledge(
            "completed",
            message="Pick-and-place mission finished",
            data={"steps": total_steps},
        )
        messenger.send_feedback(
            AgentFeedback(
                message="Mission completed successfully",
                kind="success",
                ok=True,
                progress=1.0,
                data={"completed_steps": completed},
            )
        )
        return {"ack": True, "steps": completed}

    engine.on_template(MISSION_TEMPLATE, _mission_handler)
    config.configure(engine)
    return engine, context_provider


def start_mission_mqtt_demo(
    *,
    steps: Sequence[MissionStep] = DEFAULT_STEPS,
    approval_timeout: float = 45.0,
) -> OmniLinkMQTTBridge:
    """Start an MQTT bridge exposing the mission routine command."""

    engine, context_provider = build_mission_engine(
        steps=steps, approval_timeout=approval_timeout
    )
    bridge = OmniLinkMQTTBridge(engine)
    bridge.start()

    if context_provider is not None:
        try:
            bridge.publish_context(context_provider())
        except Exception:
            # Context publishing should not prevent the demo from running.
            pass

    print(
        "Publish the command 'run pick and place mission' to",
        f"{bridge.command_topic!r} to trigger the mission workflow.",
    )
    return bridge


__all__ = [
    "DEFAULT_STEPS",
    "MISSION_TEMPLATE",
    "MissionStep",
    "build_mission_engine",
    "start_mission_mqtt_demo",
]
