"""
Real-time inference server for Neon VLA.

Runs on Thor (Jetson Orin) or any CUDA machine.
Provides HTTP API for real-time action prediction from camera + language.
"""

from __future__ import annotations

import json
import logging
import os
import time
from typing import Optional

import numpy as np
import torch

from neon.model.neon_vla import NeonVLA, NeonOutput

logger = logging.getLogger(__name__)


class NeonInferenceServer:
    """Lightweight inference server for real-time robot control.

    Loads the trained model and provides fast action prediction.
    Designed for ~50ms inference latency on Jetson Orin.
    """

    def __init__(
        self,
        model_path: str = "cagataydev/neon-g1-v1",
        device: str = "auto",
    ):
        # Select device
        if device == "auto":
            if torch.cuda.is_available():
                self.device = torch.device("cuda")
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                self.device = torch.device("mps")
            else:
                self.device = torch.device("cpu")
        else:
            self.device = torch.device(device)

        logger.info(f"Loading Neon model from {model_path} on {self.device}")
        self.model = NeonVLA.from_pretrained(model_path, load_backbone=True)
        self.model.to(self.device)
        self.model.eval()

        # Action smoothing (temporal exponential moving average)
        self._prev_actions: Optional[np.ndarray] = None
        self._smoothing_alpha = 0.7  # Higher = more responsive, lower = smoother

        logger.info(f"Neon inference server ready: {self.model}")

    @torch.no_grad()
    def predict(
        self,
        image=None,
        video_frames: list = None,
        instruction: str = "",
        proprioception: np.ndarray = None,
        smooth: bool = True,
    ) -> NeonOutput:
        """Predict next actions from observations.

        Args:
            image: Current camera frame (PIL Image)
            video_frames: Recent camera frames
            instruction: Language instruction
            proprioception: Current joint positions
            smooth: Apply temporal smoothing

        Returns:
            NeonOutput with predicted actions
        """
        start = time.time()

        output = self.model.predict(
            image=image,
            video_frames=video_frames,
            instruction=instruction,
            proprioception=proprioception,
        )

        # Temporal smoothing
        if smooth and self._prev_actions is not None:
            output.actions = (
                self._smoothing_alpha * output.actions
                + (1 - self._smoothing_alpha) * self._prev_actions
            )
            # Re-denormalize
            output.raw_actions = np.stack([
                self.model.action_space.denormalize(a) for a in output.actions
            ])

        self._prev_actions = output.actions.copy()

        elapsed = (time.time() - start) * 1000
        logger.debug(f"Inference: {elapsed:.1f}ms")

        return output

    def reset(self):
        """Reset state (call between episodes)."""
        self._prev_actions = None


def main():
    """CLI entry point for the inference server."""
    import argparse

    parser = argparse.ArgumentParser(description="Neon Inference Server")
    parser.add_argument("--model", type=str, default="cagataydev/neon-g1-v1")
    parser.add_argument("--port", type=int, default=8300)
    parser.add_argument("--host", type=str, default="0.0.0.0")
    args = parser.parse_args()

    # Create server
    server = NeonInferenceServer(model_path=args.model)

    # Simple HTTP server
    try:
        from http.server import HTTPServer, BaseHTTPRequestHandler
        import io
        from PIL import Image
        import base64

        class Handler(BaseHTTPRequestHandler):
            def do_POST(self):
                if self.path == "/predict":
                    content_len = int(self.headers.get("Content-Length", 0))
                    body = json.loads(self.rfile.read(content_len))

                    # Decode image if provided
                    image = None
                    if "image_base64" in body:
                        img_bytes = base64.b64decode(body["image_base64"])
                        image = Image.open(io.BytesIO(img_bytes))

                    # Get proprioception
                    proprio = None
                    if "proprioception" in body:
                        proprio = np.array(body["proprioception"], dtype=np.float32)

                    # Predict
                    output = server.predict(
                        image=image,
                        instruction=body.get("instruction", ""),
                        proprioception=proprio,
                    )

                    # Response
                    response = {
                        "actions": output.actions.tolist(),
                        "raw_actions": output.raw_actions.tolist() if output.raw_actions is not None else None,
                    }
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(json.dumps(response).encode())

                elif self.path == "/reset":
                    server.reset()
                    self.send_response(200)
                    self.end_headers()
                    self.wfile.write(b'{"status":"ok"}')

            def do_GET(self):
                if self.path == "/health":
                    info = {
                        "status": "ok",
                        "model": str(server.model),
                        "device": str(server.device),
                        "action_space": server.model.action_space.to_dict(),
                    }
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(json.dumps(info).encode())

        httpd = HTTPServer((args.host, args.port), Handler)
        logger.info(f"Neon inference server running on {args.host}:{args.port}")
        print(f"🤖 Neon inference server: http://{args.host}:{args.port}")
        print(f"   POST /predict — predict actions")
        print(f"   POST /reset   — reset episode state")
        print(f"   GET  /health  — server status")
        httpd.serve_forever()

    except KeyboardInterrupt:
        print("\n🤖 Neon server stopped")


if __name__ == "__main__":
    main()
