"""Pydantic-free config for the VERA policy provider.

VERA (Video-to-Embodied Robot Action, MIT/CSAIL) is a two-stage video-to-action
policy: an embodiment-agnostic **video planner** (DFoT / WAN) dreams future
frames, and an embodiment-specific **Jacobian IDM** translates the dream into
robot actions. The two stages run inside a single websocket policy server
(``vera.server.start_vera_server``); this provider is a typed client + a managed
server subprocess, mirroring the ``cosmos3`` provider's service pattern.

The config mirrors VERA's server flags 1:1 and is env-overridable so a rollout
can be driven entirely from the environment (CI / fleet) without code changes.
"""

from __future__ import annotations

import dataclasses
import os
from pathlib import Path
from typing import Any, Literal

from strands_robots.utils import (
    positive_finite_number_error,
    positive_whole_number_error,
    tcp_port_error,
)

Embodiment = Literal["pusht", "mimicgen", "allegro", "droid"]
# "mimicgen" is the working, faithful embodiment end-to-end (eef-delta -> IK
# onto a real arm). "pusht" is experimental: its server runs, but VERA's IDM
# "du" action path for pusht is not wired end-to-end upstream (documented in
# VERA's configurations/dataset/pusht.yaml), so it validates the
# provider -> server -> action plumbing rather than producing a solving
# rollout. "allegro"/"droid" are code-present but checkpoint-absent (Wave 2).

# Per-embodiment default ports (policy, viz) - match the VERA examples
# (PushT uses 8820/8821; everything else uses 8800/8801).
# Per-embodiment per-view render width the VERA WAN/DFoT planner expects. The
# server does NOT advertise this (image_resolution is None); it is the client's
# job to send each view at this width (matching VERA's own RemotePolicy
# render_size). pusht: single 252-wide view (PushTImageEnv default); mimicgen:
# 128/view (run_mimicgen_eval default); droid/allegro follow upstream eval.
_DEFAULT_RENDER_WIDTH: dict[str, int] = {
    "pusht": 252,
    "mimicgen": 128,
    "droid": 128,
    "allegro": 128,
}

_DEFAULT_PORTS: dict[str, tuple[int, int]] = {
    "pusht": (8820, 8821),
    "mimicgen": (8800, 8801),
    "allegro": (8802, 8803),
    "droid": (8804, 8805),
}


def _env(*names: str) -> str | None:
    """Return the first set (non-empty) environment variable among ``names``."""
    for n in names:
        v = os.environ.get(n)
        if v is not None and v.strip() != "":
            return v
    return None


def _env_int(name: str) -> int | None:
    v = os.environ.get(name)
    if v is None or v.strip() == "":
        return None
    try:
        return int(v)
    except ValueError:
        return None


def _viewer_port_error(value: Any, param: str, context: str) -> str | None:
    """Error text when ``value`` cannot address the MJPEG live-viewer port.

    :func:`~strands_robots.utils.tcp_port_error` owns the range and the scalar
    policy; this wrapper decides only the floor, because ``vis_port`` documents
    ``0`` as "disable the live viewer" rather than as a port to bind - the
    runner omits the ``--vis-port`` flag entirely for a falsy value. A genuine
    ``int`` zero is therefore accepted here and every other value is deferred,
    so the two ports cannot drift apart on what counts as an addressable one.

    ``bool`` is deliberately not spelled in the zero test: ``False == 0``, so a
    bare ``value == 0`` would read ``False`` as the disable spelling. The type
    identity is checked first, and a boolean falls through to
    :func:`~strands_robots.utils.tcp_port_error`, which refuses it for the same
    reason it refuses one for ``server_port``.

    Args:
        value: The caller-supplied value, after the default has been applied.
        param: The field name it came from, used in the message.
        context: Message prefix identifying the surface that received it.

    Returns:
        An error message, or ``None`` when the value is usable.
    """
    if isinstance(value, int) and not isinstance(value, bool) and value == 0:
        return None
    return tcp_port_error(value, param, context)


def _env_float(name: str) -> float | None:
    v = os.environ.get(name)
    if v is None or v.strip() == "":
        return None
    try:
        return float(v)
    except ValueError:
        return None


@dataclasses.dataclass
class VeraConfig:
    """Runtime configuration for :class:`VeraPolicy`.

    Every field maps to a VERA server flag / env var. Construction applies
    environment overrides last so callers can pin values in code and still let
    the environment win for deploy-time knobs (ports, checkpoint roots).

    Args:
        embodiment: VERA embodiment - selects the WAN/DFoT planner + Jacobian
            IDM pair and the client-side action adapter.
        host: Policy-server hostname.
        server_port: Policy-server websocket port. ``None`` applies the
            per-embodiment default; any other value must be an ``int`` in
            ``[1, 65535]``, because the client dials it and the server binds it.
        vis_port: MJPEG live-viewer port. ``None`` applies the per-embodiment
            default; ``0`` disables the viewer (the runner omits
            ``--vis-port``); any other value must be an ``int`` in
            ``[1, 65535]``.
        render_width: Per-view width, in pixels, each camera frame is resized
            to before it is sent to the planner. ``None`` applies
            ``VERA_RENDER_WIDTH`` else the per-embodiment default; any other
            value must be a positive whole number of pixels - the shared media
            domain (:func:`~strands_robots.utils.positive_whole_number_error`)
            that the recorders' ``width``/``height`` and
            :class:`~strands_robots.rendering.HybridCompositor` already use.
        algo_config: WAN planner ``algo_config.yaml`` path. Point at the omni
            config to swap the planner without retraining the IDM.
        dynamics_run_id: Jacobian/IDM checkpoint id (wandb run id); falls back
            to the embodiment's in-tree default when unset.
        text_prompt: Optional text conditioning for the video planner.
        ckpt_root: Root of the downloaded VERA checkpoints
            (``hf download sizhe-lester-li/VERA --local-dir …``). Exported to
            ``VERA_CKPT_ROOT`` for the server subprocess.
        sample_steps: WAN denoise steps (deploy default is 10; ``None`` uses the
            planner yaml's value).
        tracker_backend: IDM point tracker backend override.
        motion_plan_scale: IDM motion-plan scale override (live-tunable).
            ``None`` - the default, and what an unset ``VERA_MOTION_PLAN_SCALE``
            resolves to - leaves the server's own scale alone. Any other value
            must be a positive finite number, because it multiplies the motion
            plan the IDM turns into actions.
        teacache: Enable the near-lossless DiT teacache speedup (default True).
        teacache_thresh: teacache rel_l1 threshold (>0.15 hits a quality cliff).
        auto_launch_server: Launch + manage the server subprocess on first use.
        server_ready_timeout: Seconds to wait for the server websocket to come
            up before raising (WAN model load can be slow).
        python_executable: Interpreter used to launch the server subprocess
            (defaults to the current interpreter / ``VERA_PYTHON``).
    """

    embodiment: Embodiment = "pusht"
    host: str = "127.0.0.1"
    server_port: int | None = None
    vis_port: int | None = None
    render_width: int | None = None  # per-view pixel width sent to the server (per-embodiment default)
    algo_config: Path | None = None
    dynamics_run_id: str | None = None
    text_prompt: str | None = None
    ckpt_root: Path | None = None
    wan_ckpt_root: Path | None = None  # frozen Wan2.1-T2V-1.3B base (mimicgen/omni); env VERA_WAN_CKPT_ROOT
    sample_steps: int | None = None
    tracker_backend: str | None = None
    motion_plan_scale: float | None = None
    teacache: bool = True
    teacache_thresh: float = 0.10
    auto_launch_server: bool = True
    server_ready_timeout: float = 600.0
    python_executable: str | None = None
    # --- server launch mode -------------------------------------------------
    server_mode: str = "subprocess"  # "subprocess" | "docker"
    docker_image: str = "strands-vera-server:latest"
    docker_container_name: str | None = None  # default: vera-server-<embodiment>
    docker_gpus: str = "all"  # --gpus value (e.g. "all" or "device=0")
    docker_extra_args: list[str] | None = None  # extra `docker run` args (list, no shell)

    def __post_init__(self) -> None:
        # Apply per-embodiment port defaults when not explicitly set.
        default_policy, default_vis = _DEFAULT_PORTS.get(self.embodiment, (8800, 8801))
        if self.server_port is None:
            self.server_port = _env_int("VERA_SERVER_PORT") or default_policy
        if self.vis_port is None:
            env_vis = _env_int("VERA_VIS_PORT")
            self.vis_port = env_vis if env_vis is not None else default_vis

        # Both ports are validated here, on the effective value, because this is
        # the one funnel every caller passes through - the ``VeraPolicy``
        # keywords, a pre-built config handed to it, and the ``VERA_*_PORT``
        # environment overrides above - and because a port reaches three
        # consumers under three different coercions: the provider dials
        # ``int(server_port or 0)``, :attr:`server_uri` interpolates the field
        # verbatim, and the runner's argv carries ``str(server_port)``. An
        # unusable value is therefore not merely refused late; it is applied as
        # three different ports (``2.7`` dials ``:2``, reports
        # ``ws://host:2.7`` and launches ``--port 2.7``, so the client cannot
        # reach the server it just started). Refusing before any client or
        # runner is built leaves nothing half-configured behind.
        if (err := tcp_port_error(self.server_port, "server_port", type(self).__name__)) is not None:
            raise ValueError(err)
        if (err := _viewer_port_error(self.vis_port, "vis_port", type(self).__name__)) is not None:
            raise ValueError(err)

        # ``render_width`` is a pixel count, so it takes the shared media domain
        # rather than a local rule: it is the same quantity as the recorders'
        # ``width``/``height`` and ``HybridCompositor.default_width``, and
        # ``positive_whole_number_error`` names pixels as one of the two families
        # it exists for. Applied here, on the effective value, for the reason the
        # ports above are - ``render_width`` is read only inside
        # ``_extract_frame``, which runs per frame *after* ``_ensure_started``
        # has launched the WAN server subprocess and completed the handshake, so
        # an unusable width surfaced there costs a model load (up to
        # ``server_ready_timeout``) before reporting a value that was wrong at
        # construction.
        #
        # The env override is read with ``is not None``, matching ``vis_port``
        # above rather than the ``or`` this line used to carry. The two spellings
        # are not interchangeable here: ``VERA_RENDER_WIDTH=0`` is falsy, so the
        # override was discarded and the per-embodiment default silently applied
        # in its place - a width of 0 cannot be honored, and the caller who asked
        # for it is owed the refusal below, not 128 under a success.
        if self.render_width is None:
            env_width = _env_int("VERA_RENDER_WIDTH")
            self.render_width = env_width if env_width is not None else _DEFAULT_RENDER_WIDTH.get(self.embodiment, 128)
        if (err := positive_whole_number_error(self.render_width, "render_width", type(self).__name__)) is not None:
            raise ValueError(err)
        # Normalized to a plain ``int`` here because the domain accepts any real
        # scalar with an integral value - a ``128.0`` or a ``np.int64`` passes -
        # and the consumer requires a true ``int``: ``_resize_frame`` compares it
        # against ``frame.shape`` and hands it to ``Image.resize``, and it is
        # sent to the server as ``view_widths``. This is the normalization the
        # shared domain documents as the caller's obligation, and it is why the
        # ``int()`` at each read site is no longer needed.
        self.render_width = int(self.render_width)

        # Environment overrides (deploy/CI win over code defaults).
        if self.algo_config is None:
            ac = _env("VERA_ALGO_CONFIG")
            self.algo_config = Path(ac) if ac else None
        if self.dynamics_run_id is None:
            self.dynamics_run_id = _env("VERA_DYNAMICS_RUN_ID")
        if self.text_prompt is None:
            self.text_prompt = _env("VERA_TEXT_PROMPT")
        if self.ckpt_root is None:
            cr = _env(
                "VERA_CKPT_ROOT",
                f"VERA_{self.embodiment.upper()}_CKPT_ROOT",
                "VERA_MIMICGEN_CKPT_DIR" if self.embodiment == "mimicgen" else "",
            )
            self.ckpt_root = Path(cr) if cr else None
        if self.wan_ckpt_root is None:
            wr = _env("VERA_WAN_CKPT_ROOT")
            self.wan_ckpt_root = Path(wr) if wr else None
        if self.sample_steps is None:
            self.sample_steps = _env_int("VERA_SAMPLE_STEPS")
        if self.tracker_backend is None:
            self.tracker_backend = _env("VERA_TRACKER_BACKEND")
        if self.motion_plan_scale is None:
            self.motion_plan_scale = _env_float("VERA_MOTION_PLAN_SCALE")
        # Checked here, on the effective value, for the reason the two ports and
        # ``render_width`` above are: this is the one funnel every caller passes
        # through, and the only place the ``VERA_MOTION_PLAN_SCALE`` override just
        # applied can still be refused. ``_env_float`` returns whatever ``float()``
        # accepts, so ``nan``, ``inf``, ``1e999`` and a negative all reach the field
        # from the environment; the field was checked nowhere, so a ``str`` or a
        # ``list`` reached it from a keyword too.
        #
        # It is the third scale in this package and takes the domain the other two
        # already take (``translation_scale``, ``ik_smoothing``): a multiplier on a
        # motion plan has no usable non-positive or non-finite value. ``None`` stays
        # valid because it is the documented opt-out - it gates the ``configure``
        # call away entirely, so "leave the server's scale alone" and "scale the
        # plan to nothing" stay different requests.
        #
        # Leaving it to be refused downstream is not an option, because nothing
        # downstream refuses it: ``_ensure_started`` sends the value with
        # ``self._client.configure(...)`` inside a best-effort ``except Exception``
        # that logs at INFO and marks the policy started regardless. A value
        # ``float()`` cannot convert is therefore neither applied nor reported -
        # the rollout proceeds at whatever scale the server already had while the
        # config says otherwise.
        if self.motion_plan_scale is not None:
            if (
                err := positive_finite_number_error(self.motion_plan_scale, "motion_plan_scale", type(self).__name__)
            ) is not None:
                raise ValueError(err)
            # Normalized for the reason ``render_width`` is: the shared domain admits
            # any real scalar, so an ``int`` 1 or a ``np.float64`` passes it, while
            # the field is declared ``float``.
            self.motion_plan_scale = float(self.motion_plan_scale)
        if self.python_executable is None:
            self.python_executable = _env("VERA_PYTHON")
        _sm = _env("VERA_SERVER_MODE")
        if _sm:
            self.server_mode = _sm
        _di = _env("VERA_DOCKER_IMAGE")
        if _di:
            self.docker_image = _di
        _dg = _env("VERA_DOCKER_GPUS")
        if _dg:
            self.docker_gpus = _dg
        if self.docker_container_name is None:
            self.docker_container_name = _env("VERA_DOCKER_CONTAINER") or f"vera-server-{self.embodiment}"

        # Coerce string paths to Path (defensive - callers may pass str).
        if self.algo_config is not None and not isinstance(self.algo_config, Path):
            self.algo_config = Path(self.algo_config)
        if self.ckpt_root is not None and not isinstance(self.ckpt_root, Path):
            self.ckpt_root = Path(self.ckpt_root)
        if self.wan_ckpt_root is not None and not isinstance(self.wan_ckpt_root, Path):
            self.wan_ckpt_root = Path(self.wan_ckpt_root)

    @property
    def server_uri(self) -> str:
        """Websocket URI the client connects to (``ws://<host>:<server_port>``)."""
        return f"ws://{self.host}:{self.server_port}"

    def server_env(self) -> dict[str, str]:
        """Environment overlay for the server subprocess (checkpoints, tracker)."""
        env: dict[str, str] = {}
        if self.ckpt_root is not None:
            env["VERA_CKPT_ROOT"] = str(self.ckpt_root)
        if self.wan_ckpt_root is not None:
            env["VERA_WAN_CKPT_ROOT"] = str(self.wan_ckpt_root)
        if self.tracker_backend is not None:
            env["VERA_TRACKER_BACKEND"] = self.tracker_backend
        if self.dynamics_run_id is not None:
            env["VERA_DYNAMICS_RUN_ID"] = str(self.dynamics_run_id)
        return env
