"""
SO-ARM101 RCAN config generator.

Generates a ready-to-use .rcan.yaml from detected ports and arm config.

Usage:
    from castor.hardware.so_arm101.config_generator import generate_config
    yaml_str = generate_config(follower_port="/dev/ttyACM0", robot_name="my_arm")
"""

from __future__ import annotations

import datetime
import os

from castor.hardware.so_arm101.constants import DEFAULT_BAUD, FOLLOWER_MOTORS, LEADER_MOTORS


def generate_config(
    robot_name: str = "so_arm101",
    follower_port: str = "/dev/ttyACM0",
    leader_port: str | None = None,
    baud: int = DEFAULT_BAUD,
    provider: str = "ollama",
    model: str = "llama3.2:3b",
    rrn: str | None = None,
) -> str:
    """
    Generate a complete RCAN config YAML string for the SO-ARM101.

    Args:
        robot_name: Display name for the robot.
        follower_port: Serial port for the follower arm.
        leader_port: Serial port for the leader arm (optional, for bimanual).
        baud: Baudrate (default 1,000,000).
        provider: LLM provider name.
        model: LLM model ID.
        rrn: Robot Registration Number (optional).
    """
    follower_joint_names = [m["joint"] for m in FOLLOWER_MOTORS]
    follower_ids = [m["id"] for m in FOLLOWER_MOTORS]

    leader_section = ""
    if leader_port:
        leader_joint_names = [m["joint"] for m in LEADER_MOTORS]
        leader_ids = [m["id"] for m in LEADER_MOTORS]
        leader_section = f"""
  - id: leader_arm
    protocol: feetech
    port: {leader_port}
    baud: {baud}
    servo_ids: {leader_ids}
    joint_names: {leader_joint_names}
    operating_mode: position
    role: leader  # Teleoperation leader; mirrors pose to follower_arm
"""

    rrn_line = f"rrn: {rrn}" if rrn else "# rrn: RRN-000000000000  # assign via: castor register"
    bimanual_note = "(bimanual)" if leader_port else "(single-arm follower)"
    date = datetime.date.today().isoformat()

    return f"""# OpenCastor RCAN config — SO-ARM101 {bimanual_note}
# Generated by: castor arm setup
# Date: {date}
# Docs: https://opencastor.com/docs
# SO-ARM101: https://huggingface.co/docs/lerobot/so101

robot_name: "{robot_name}"
rcan_version: "1.4"
{rrn_line}

# ── Hardware ──────────────────────────────────────────────────────────────────
platform: so_arm101
robot_type: manipulator

drivers:
  - id: follower_arm
    protocol: feetech
    port: {follower_port}
    baud: {baud}
    servo_ids: {follower_ids}
    joint_names: {follower_joint_names}
    operating_mode: position
    role: follower
{leader_section}
# ── AI Brain ──────────────────────────────────────────────────────────────────
provider: {provider}
model: "{model}"

task_routing:
  sensor_poll: sensor          # Fast sensor reads → no LLM
  safety: planner              # Safety always uses full planner
  grasp: planner               # Grasp planning → LLM
  trajectory: planner
  default: sensor

# ── Safety ────────────────────────────────────────────────────────────────────
safety:
  joint_limits:
    shoulder_pan:  [-2.09, 2.09]    # ±120°
    shoulder_lift: [-1.57, 1.57]    # ±90°
    elbow_flex:    [-2.44, 2.44]    # ±140°
    wrist_flex:    [-1.57, 1.57]    # ±90°
    wrist_roll:    [-3.14, 3.14]    # ±180°
    gripper:       [0.0,   1.0]     # 0 = open, 1 = closed
  max_velocity_rad_s: 1.5
  collision_detection: true
  emergency_stop_on_overload: true

# ── Teleoperation ─────────────────────────────────────────────────────────────
# Uncomment to enable real-time leader→follower mirroring
# teleop:
#   enabled: true
#   leader_driver: leader_arm
#   follower_driver: follower_arm
#   scale: 1.0

# ── Recording / LeRobot integration ──────────────────────────────────────────
# Uncomment to record episodes for imitation learning
# episodes:
#   enabled: false
#   output_dir: ~/so_arm101_episodes
#   format: lerobot   # compatible with HuggingFace LeRobot
#   fps: 30
"""


def write_config(path: str, **kwargs) -> str:
    """Write a generated RCAN config to a file. Returns the path."""
    yaml_str = generate_config(**kwargs)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w") as f:
        f.write(yaml_str)
    return path
