"""
Neon Tool — Strands Agent tool for controlling the G1 via natural language.

Usage:
    from strands import Agent
    from neon.tools import neon_tool

    agent = Agent(tools=[neon_tool])
    agent("Neon, pick up the red cup and bring it to me")
"""

from __future__ import annotations

import logging
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)

# Lazy-loaded globals
_server = None
_controller = None


def _get_server():
    """Lazy-load the inference server."""
    global _server
    if _server is None:
        import os
        from neon.inference.server import NeonInferenceServer

        model_path = os.getenv("NEON_MODEL_PATH", "cagataydev/neon-g1-v1")
        _server = NeonInferenceServer(model_path=model_path)
    return _server


def neon_tool(
    action: str = "predict",
    instruction: str = "",
    image_path: Optional[str] = None,
    model_path: str = "cagataydev/neon-g1-v1",
    robot_ip: str = "192.168.123.10",
    control_mode: str = "arms_only",
    max_steps: int = 200,
) -> Dict[str, Any]:
    """
    🤖 Neon — Control the G1 humanoid with natural language.

    Predict robot actions from camera images and language instructions,
    or run a full control loop on the physical robot.

    Args:
        action: Action to perform:
            - "predict": Predict actions from image + instruction
            - "control": Run closed-loop control on G1
            - "status": Show model and robot status
            - "reset": Reset episode state
        instruction: Natural language instruction (e.g., "pick up the red cup")
        image_path: Path to camera image (for predict action)
        model_path: HuggingFace model ID or local path
        robot_ip: G1 robot IP address
        control_mode: "arms_only", "upper_body", or "whole_body"
        max_steps: Max control steps (for control action)

    Returns:
        Dict with predicted actions or control status

    Examples:
        # Predict actions from image
        neon_tool(action="predict", instruction="pick up the cup", image_path="/tmp/camera.jpg")

        # Run control loop on G1
        neon_tool(action="control", instruction="pick up the cup", robot_ip="192.168.123.10")

        # Check status
        neon_tool(action="status")
    """
    try:
        if action == "predict":
            server = _get_server()

            # Load image
            image = None
            if image_path:
                from PIL import Image

                image = Image.open(image_path)

            # Predict
            output = server.predict(image=image, instruction=instruction)

            return {
                "status": "success",
                "content": [
                    {
                        "text": (
                            f"🤖 Neon prediction:\n"
                            f"  Instruction: {instruction}\n"
                            f"  Action shape: {output.actions.shape}\n"
                            f"  Upper body: {output.upper_body.shape if output.upper_body is not None else 'N/A'}\n"
                            f"  Locomotion: {output.locomotion[0].tolist() if output.locomotion is not None else 'N/A'}\n"
                            f"  First action: {output.actions[0, :5].tolist()}..."
                        )
                    }
                ],
                "actions": output.actions.tolist(),
                "raw_actions": output.raw_actions.tolist() if output.raw_actions is not None else None,
            }

        elif action == "control":
            from neon.inference.g1_controller import G1Config, G1Controller

            server = _get_server()
            config = G1Config(robot_ip=robot_ip, mode=control_mode)
            controller = G1Controller(config)
            controller.connect()

            if not controller.is_connected():
                return {
                    "status": "error",
                    "content": [{"text": f"🤖 Failed to connect to G1 at {robot_ip}"}],
                }

            try:
                controller.run_control_loop(server, instruction=instruction, max_steps=max_steps)
                return {
                    "status": "success",
                    "content": [{"text": f"🤖 Control loop completed ({max_steps} steps)"}],
                }
            finally:
                controller.disconnect()

        elif action == "status":
            info = {
                "model_loaded": _server is not None,
                "model_path": model_path,
            }
            if _server is not None:
                info["model_info"] = str(_server.model)
                info["device"] = str(_server.device)
                info["action_space"] = _server.model.action_space.to_dict()

            return {
                "status": "success",
                "content": [{"text": f"🤖 Neon status:\n{info}"}],
            }

        elif action == "reset":
            if _server is not None:
                _server.reset()
            return {
                "status": "success",
                "content": [{"text": "🤖 Neon episode state reset"}],
            }

        else:
            return {
                "status": "error",
                "content": [{"text": f"Unknown action: {action}. Use: predict, control, status, reset"}],
            }

    except Exception as e:
        logger.error(f"Neon tool error: {e}")
        return {"status": "error", "content": [{"text": f"🤖 Error: {e}"}]}
