"""
HOLD → glTF Exporter

Exports inference pause data as universal 3D meshes.
Import into Unity, Unreal, Godot, Blender - anything.

Mesh structure:
- Each action candidate = a 3D trajectory tube
- Tube radius = probability (thicker = more likely)
- Tube color = expected value (red=bad, green=good)
- Trajectory path = imagined future states projected to 3D
"""

import json
import struct
import base64
import math
from typing import Dict, Any, List, Optional, Tuple
from dataclasses import dataclass
import numpy as np


@dataclass
class TrajectoryMesh:
    """A single action trajectory as mesh data."""
    action_idx: int
    action_name: str
    probability: float
    expected_value: float
    path_points: List[Tuple[float, float, float]]  # 3D trajectory
    

def _numpy_to_list(obj):
    """Recursively convert numpy arrays to lists."""
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, (np.float32, np.float64)):
        return float(obj)
    elif isinstance(obj, (np.int32, np.int64)):
        return int(obj)
    elif isinstance(obj, dict):
        return {k: _numpy_to_list(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_numpy_to_list(v) for v in obj]
    return obj


def hold_to_trajectories(hold_data: Dict[str, Any]) -> List[TrajectoryMesh]:
    """
    Convert HOLD data to trajectory meshes.
    
    Extracts action probabilities and imagination rollouts,
    maps them to 3D space for visualization.
    """
    trajectories = []
    
    # Get action probabilities
    action_probs = hold_data.get("action_probs", [])
    if isinstance(action_probs, np.ndarray):
        action_probs = action_probs.tolist()
    
    action_names = hold_data.get("action_names", [])
    if not action_names:
        action_names = [f"action_{i}" for i in range(len(action_probs))]
    
    # Get imagination data if available
    imagination = hold_data.get("imagination", {})
    trajectories_data = imagination.get("trajectories", [])
    
    # Get expected values
    expected_values = hold_data.get("expected_values", [0.0] * len(action_probs))
    if isinstance(expected_values, np.ndarray):
        expected_values = expected_values.tolist()
    
    # If no trajectory data, generate from probabilities alone
    if not trajectories_data:
        trajectories_data = _generate_synthetic_paths(action_probs, expected_values)
    
    # Build trajectory meshes
    for i, prob in enumerate(action_probs):
        ev = expected_values[i] if i < len(expected_values) else 0.0
        
        # Get path for this action
        if i < len(trajectories_data):
            path = trajectories_data[i]
        else:
            # Generate radial path based on action index
            angle = (2 * math.pi * i) / len(action_probs)
            path = _generate_radial_path(angle, prob, ev)
        
        trajectories.append(TrajectoryMesh(
            action_idx=i,
            action_name=action_names[i] if i < len(action_names) else f"action_{i}",
            probability=float(prob),
            expected_value=float(ev),
            path_points=path
        ))
    
    return trajectories


def _generate_synthetic_paths(probs: List[float], evs: List[float]) -> List[List[Tuple[float, float, float]]]:
    """Generate 3D paths when no imagination data available."""
    paths = []
    n_actions = len(probs)
    
    for i, (prob, ev) in enumerate(zip(probs, evs)):
        angle = (2 * math.pi * i) / n_actions
        paths.append(_generate_radial_path(angle, prob, ev))
    
    return paths


def _generate_radial_path(angle: float, prob: float, ev: float, steps: int = 10) -> List[Tuple[float, float, float]]:
    """Generate a single radial trajectory path."""
    path = []
    
    # Path extends outward from origin
    # Length proportional to probability
    # Height based on expected value
    max_length = 2.0 + prob * 3.0  # 2-5 units based on prob
    
    for step in range(steps + 1):
        t = step / steps
        
        # Position along radial line
        x = math.cos(angle) * t * max_length
        z = math.sin(angle) * t * max_length
        
        # Height based on EV, with slight curve
        y = ev * t + 0.5 * math.sin(t * math.pi) * abs(ev)
        
        path.append((x, y, z))
    
    return path


def _create_tube_mesh(path: List[Tuple[float, float, float]], radius: float, segments: int = 8) -> Tuple[List[float], List[int]]:
    """
    Create a tube mesh along a path.
    
    Returns (vertices, indices) for the mesh.
    Vertices are [x,y,z, x,y,z, ...] flat list.
    """
    vertices = []
    indices = []
    
    if len(path) < 2:
        return vertices, indices
    
    # Generate rings along the path
    for i, point in enumerate(path):
        # Calculate tangent direction
        if i == 0:
            tangent = _normalize(_subtract(path[1], path[0]))
        elif i == len(path) - 1:
            tangent = _normalize(_subtract(path[-1], path[-2]))
        else:
            tangent = _normalize(_subtract(path[i+1], path[i-1]))
        
        # Find perpendicular vectors
        up = (0, 1, 0) if abs(tangent[1]) < 0.9 else (1, 0, 0)
        right = _normalize(_cross(tangent, up))
        actual_up = _cross(right, tangent)
        
        # Create ring of vertices
        for seg in range(segments):
            angle = (2 * math.pi * seg) / segments
            
            # Offset from center
            offset_x = math.cos(angle) * radius * right[0] + math.sin(angle) * radius * actual_up[0]
            offset_y = math.cos(angle) * radius * right[1] + math.sin(angle) * radius * actual_up[1]
            offset_z = math.cos(angle) * radius * right[2] + math.sin(angle) * radius * actual_up[2]
            
            vertices.extend([
                point[0] + offset_x,
                point[1] + offset_y,
                point[2] + offset_z
            ])
    
    # Create triangle indices connecting rings
    for ring in range(len(path) - 1):
        for seg in range(segments):
            next_seg = (seg + 1) % segments
            
            # Current ring indices
            curr_a = ring * segments + seg
            curr_b = ring * segments + next_seg
            
            # Next ring indices
            next_a = (ring + 1) * segments + seg
            next_b = (ring + 1) * segments + next_seg
            
            # Two triangles per quad
            indices.extend([curr_a, next_a, curr_b])
            indices.extend([curr_b, next_a, next_b])
    
    return vertices, indices


def _normalize(v: Tuple[float, float, float]) -> Tuple[float, float, float]:
    length = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    if length < 1e-8:
        return (0, 1, 0)
    return (v[0]/length, v[1]/length, v[2]/length)


def _subtract(a: Tuple[float, float, float], b: Tuple[float, float, float]) -> Tuple[float, float, float]:
    return (a[0]-b[0], a[1]-b[1], a[2]-b[2])


def _cross(a: Tuple[float, float, float], b: Tuple[float, float, float]) -> Tuple[float, float, float]:
    return (
        a[1]*b[2] - a[2]*b[1],
        a[2]*b[0] - a[0]*b[2],
        a[0]*b[1] - a[1]*b[0]
    )


def _ev_to_color(ev: float) -> Tuple[float, float, float, float]:
    """Convert expected value to RGBA color."""
    # Normalize EV to [-1, 1] range roughly
    ev_norm = max(-1, min(1, ev))
    
    if ev_norm >= 0:
        # Positive: green
        return (0.2, 0.5 + 0.5 * ev_norm, 0.2, 1.0)
    else:
        # Negative: red
        return (0.5 - 0.5 * ev_norm, 0.2, 0.2, 1.0)


def export_hold_to_gltf(hold_data: Dict[str, Any], output_path: str, binary: bool = True) -> str:
    """
    Export HOLD data to glTF/GLB file.
    
    Args:
        hold_data: HOLD snapshot with action_probs, imagination, etc.
        output_path: Where to save the file
        binary: If True, saves .glb (single binary). If False, saves .gltf (JSON).
    
    Returns:
        Path to the saved file
    """
    hold_data = _numpy_to_list(hold_data)
    trajectories = hold_to_trajectories(hold_data)
    
    # Build glTF structure
    gltf = {
        "asset": {
            "version": "2.0",
            "generator": "cascade-hold-exporter",
            "extras": {
                "hold_timestamp": hold_data.get("timestamp", ""),
                "model_card": hold_data.get("model_card", {}),
                "decision_id": hold_data.get("decision_id", "")
            }
        },
        "scene": 0,
        "scenes": [{"nodes": list(range(len(trajectories)))}],
        "nodes": [],
        "meshes": [],
        "accessors": [],
        "bufferViews": [],
        "buffers": [],
        "materials": []
    }
    
    # Collect all binary data
    all_binary = bytearray()
    
    for i, traj in enumerate(trajectories):
        # Create material for this trajectory (color by EV)
        color = _ev_to_color(traj.expected_value)
        gltf["materials"].append({
            "name": f"mat_{traj.action_name}",
            "pbrMetallicRoughness": {
                "baseColorFactor": list(color),
                "metallicFactor": 0.1,
                "roughnessFactor": 0.8
            },
            "extras": {
                "action_idx": traj.action_idx,
                "action_name": traj.action_name,
                "probability": traj.probability,
                "expected_value": traj.expected_value
            }
        })
        
        # Create tube mesh
        # Radius based on probability (0.05 to 0.3)
        radius = 0.05 + traj.probability * 0.25
        vertices, indices = _create_tube_mesh(traj.path_points, radius)
        
        if not vertices or not indices:
            continue
        
        # Pack vertex data
        vertex_data = struct.pack(f'{len(vertices)}f', *vertices)
        vertex_offset = len(all_binary)
        all_binary.extend(vertex_data)
        
        # Pad to 4-byte boundary
        while len(all_binary) % 4 != 0:
            all_binary.append(0)
        
        # Pack index data
        index_data = struct.pack(f'{len(indices)}H', *indices)
        index_offset = len(all_binary)
        all_binary.extend(index_data)
        
        # Pad to 4-byte boundary
        while len(all_binary) % 4 != 0:
            all_binary.append(0)
        
        # Calculate bounds
        n_verts = len(vertices) // 3
        xs = [vertices[j*3] for j in range(n_verts)]
        ys = [vertices[j*3+1] for j in range(n_verts)]
        zs = [vertices[j*3+2] for j in range(n_verts)]
        
        # Buffer views
        vertex_view_idx = len(gltf["bufferViews"])
        gltf["bufferViews"].append({
            "buffer": 0,
            "byteOffset": vertex_offset,
            "byteLength": len(vertex_data),
            "target": 34962  # ARRAY_BUFFER
        })
        
        index_view_idx = len(gltf["bufferViews"])
        gltf["bufferViews"].append({
            "buffer": 0,
            "byteOffset": index_offset,
            "byteLength": len(index_data),
            "target": 34963  # ELEMENT_ARRAY_BUFFER
        })
        
        # Accessors
        position_accessor_idx = len(gltf["accessors"])
        gltf["accessors"].append({
            "bufferView": vertex_view_idx,
            "byteOffset": 0,
            "componentType": 5126,  # FLOAT
            "count": n_verts,
            "type": "VEC3",
            "min": [min(xs), min(ys), min(zs)],
            "max": [max(xs), max(ys), max(zs)]
        })
        
        index_accessor_idx = len(gltf["accessors"])
        gltf["accessors"].append({
            "bufferView": index_view_idx,
            "byteOffset": 0,
            "componentType": 5123,  # UNSIGNED_SHORT
            "count": len(indices),
            "type": "SCALAR"
        })
        
        # Mesh
        mesh_idx = len(gltf["meshes"])
        gltf["meshes"].append({
            "name": f"trajectory_{traj.action_name}",
            "primitives": [{
                "attributes": {"POSITION": position_accessor_idx},
                "indices": index_accessor_idx,
                "material": i,
                "mode": 4  # TRIANGLES
            }]
        })
        
        # Node
        gltf["nodes"].append({
            "name": traj.action_name,
            "mesh": mesh_idx,
            "extras": {
                "probability": traj.probability,
                "expected_value": traj.expected_value
            }
        })
    
    # Add buffer
    gltf["buffers"].append({
        "byteLength": len(all_binary)
    })
    
    # Save file
    if binary:
        # GLB format
        if not output_path.endswith('.glb'):
            output_path = output_path.rsplit('.', 1)[0] + '.glb'
        
        json_str = json.dumps(gltf)
        json_bytes = json_str.encode('utf-8')
        
        # Pad JSON to 4-byte boundary
        while len(json_bytes) % 4 != 0:
            json_bytes += b' '
        
        # GLB header
        glb_header = struct.pack('<4sII', b'glTF', 2, 12 + 8 + len(json_bytes) + 8 + len(all_binary))
        
        # JSON chunk
        json_chunk = struct.pack('<II', len(json_bytes), 0x4E4F534A) + json_bytes
        
        # BIN chunk
        bin_chunk = struct.pack('<II', len(all_binary), 0x004E4942) + bytes(all_binary)
        
        with open(output_path, 'wb') as f:
            f.write(glb_header + json_chunk + bin_chunk)
    else:
        # glTF format with embedded base64
        if not output_path.endswith('.gltf'):
            output_path = output_path.rsplit('.', 1)[0] + '.gltf'
        
        gltf["buffers"][0]["uri"] = "data:application/octet-stream;base64," + base64.b64encode(bytes(all_binary)).decode('ascii')
        
        with open(output_path, 'w') as f:
            json.dump(gltf, f, indent=2)
    
    return output_path


class HoldGLTFExporter:
    """
    Streaming exporter - accumulates HOLD snapshots into animated glTF.
    
    Usage:
        exporter = HoldGLTFExporter()
        
        # During inference
        exporter.add_frame(hold_data, timestamp=0.0)
        exporter.add_frame(hold_data, timestamp=0.1)
        ...
        
        # Save animation
        exporter.export("decision_sequence.glb")
    """
    
    def __init__(self):
        self.frames: List[Tuple[float, Dict[str, Any]]] = []
    
    def add_frame(self, hold_data: Dict[str, Any], timestamp: float = None):
        """Add a HOLD snapshot as animation frame."""
        if timestamp is None:
            timestamp = len(self.frames) * 0.1
        self.frames.append((timestamp, _numpy_to_list(hold_data)))
    
    def export(self, output_path: str) -> str:
        """Export all frames as animated glTF."""
        if not self.frames:
            raise ValueError("No frames to export")
        
        if len(self.frames) == 1:
            # Single frame - just export static
            return export_hold_to_gltf(self.frames[0][1], output_path)
        
        # For animated export, we merge all frame data
        # Each action gets keyframed scale (probability) and color (EV)
        # This is a simplified version - full animation would need more work
        
        # For now, export the last frame with all frame data in extras
        final_hold = self.frames[-1][1].copy()
        final_hold["animation_frames"] = [
            {"timestamp": t, "action_probs": f.get("action_probs", [])}
            for t, f in self.frames
        ]
        
        return export_hold_to_gltf(final_hold, output_path)
    
    def clear(self):
        """Clear accumulated frames."""
        self.frames.clear()


# Convenience function
def export_hold(hold_data: Dict[str, Any], path: str = "hold_export.glb") -> str:
    """Quick export of HOLD data to glTF."""
    return export_hold_to_gltf(hold_data, path)


if __name__ == "__main__":
    # Demo export
    demo_hold = {
        "action_probs": [0.4, 0.25, 0.15, 0.1, 0.05, 0.03, 0.01, 0.01],
        "action_names": ["forward", "left", "right", "back", "jump", "attack", "use", "wait"],
        "expected_values": [0.8, 0.3, 0.2, -0.1, 0.5, -0.3, 0.0, -0.5],
        "model_card": {
            "name": "DreamerV3",
            "latent_dim": 5120
        }
    }
    
    path = export_hold(demo_hold, "demo_hold.glb")
    print(f"Exported to: {path}")
    print("Import this file into Unity, Unreal, Godot, or Blender!")
