import os
import sys
import time
import numpy as np
import json
from pathlib import Path

# Add project root and ISOLATE paths
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))
sys.path.append(str(project_root / 'scripts' / 'ISOLATE'))

# Import the model wrapper
try:
    from scripts.ISOLATE.model_wrapper import ShenronRadarModel
except ImportError as e:
    print(f"Error: Failed to import ShenronRadarModel. Ensure scripts/ISOLATE/model_wrapper.py exists. ({e})")
    sys.exit(1)

def _get_gpu_info():
    """Retrieve GPU hardware info for telemetry display."""
    try:
        import torch
        if not torch.cuda.is_available():
            return {"name": "CPU Fallback", "vram_gb": 0, "backend": "CPU"}
            
        # Try to get device name safely
        try:
            name = torch.cuda.get_device_name(0)
        except:
            name = "NVIDIA Device"
            
        # Try to get properties safely
        try:
            props = torch.cuda.get_device_properties(0)
            vram_total = props.total_memory / (1024**3)
        except:
            vram_total = 0
            
        return {"name": name, "vram_gb": round(vram_total, 1), "backend": "CUDA"}
    except Exception as e:
        return {"name": f"Detection Error ({type(e).__name__})", "vram_gb": 0, "backend": "Unknown"}

def process_session(session_path):
    print(f"\n>>> Processing session: {session_path.name}")
    
    lidar_dir = session_path / "lidar"
    if not lidar_dir.exists():
        print(f"    [SKIP] No 'lidar' folder found.")
        return
        
    # Find all .npy files in lidar/
    lidar_files = sorted(list(lidar_dir.glob("*.npy")))
    if not lidar_files:
        print(f"    [SKIP] No .npy files in 'lidar' folder.")
        return

    from scripts.ISOLATE.shenron_orchestrator import ShenronOrchestrator
    
    orchestrator = ShenronOrchestrator(radar_types=['awrl1432', 'radarbook', 'ti_cascade'])
    
    # -----------------------------------------------------------------------
    # TELEMETRY: Init Phase
    # -----------------------------------------------------------------------
    print(f"    [DIAGNOSTIC] Step 1: Initializing models...", flush=True)
    radar_specs = orchestrator.init_models(session_path)
    
    print(f"    [DIAGNOSTIC] Step 2: Collecting metadata...", flush=True)
    gpu_info = _get_gpu_info()
    
    telemetry_init = {
        "gpu": gpu_info,
        "radars": radar_specs,
        "total_frames": len(lidar_files),
        "session": session_path.name,
    }
    print(f"[SHENRON_INIT]{json.dumps(telemetry_init)}", flush=True)
    
    # -----------------------------------------------------------------------
    # MAIN PROCESSING LOOP
    # -----------------------------------------------------------------------
    print(f"    [DIAGNOSTIC] Step 3: Starting main loop for {len(lidar_files)} frames...", flush=True)
    total_frames = len(lidar_files)
    frame_times = []
    
    for frame_idx, lidar_file in enumerate(lidar_files):
        frame_start = time.time()
        
        # Stop if requested
        if orchestrator.check_stop_flag():
            break

        # Process through the unified orchestrator
        frame_results = orchestrator.process_frame(lidar_file, session_path, save_adc=True)
        
        if frame_results is None:
            continue
            
        # Timing
        frame_elapsed = time.time() - frame_start
        frame_times.append(frame_elapsed)
        
        avg_time = sum(frame_times[-10:]) / len(frame_times[-10:])
        remaining = total_frames - (frame_idx + 1)
        eta_seconds = remaining * avg_time
        
        eta_str = f"{int(eta_seconds // 60)}m {int(eta_seconds % 60)}s" if eta_seconds > 60 else f"{int(eta_seconds)}s"
        progress_pct = int(((frame_idx + 1) / total_frames) * 100)
        
        telemetry_frame = {
            "frame": frame_idx + 1,
            "total": total_frames,
            "pct": progress_pct,
            "fps": round(1.0 / frame_elapsed, 2) if frame_elapsed > 0 else 0,
            "elapsed": round(frame_elapsed, 2),
            "eta": eta_str,
            "metrics": {}
        }
        
        for r_type, m in frame_results.items():
            telemetry_frame["metrics"][r_type] = {
                "snr": round(m.get("peak_snr_db", 0), 1),
                "pts": m.get("pts", 0), 
                "peak": round(m.get("peak_magnitude", 0), 1),
                "bins": m.get("active_bins", 0),
                "az_std": round(m.get("azimuth_variance", 0), 4),
                "spread": round(m.get("peak_azimuth_spread_deg", 0), 1),
            }
        
        print(f"[SHENRON_STEP]{json.dumps(telemetry_frame)}", flush=True)

def main():
    data_root = project_root / "data"
    if not data_root.exists():
        print(f"Error: {data_root} not found.")
        return
        
    sessions = sorted([d for d in data_root.iterdir() if d.is_dir()])
    
    if not sessions:
        print("No simulation sessions found in data/.")
        return
        
    print(f"Found {len(sessions)} sessions.")
    
    for session in sessions:
        if (session / "frames.jsonl").exists():
            process_session(session)
            
    print("\n" + "="*50)
    print("SHENRON BATCH PROCESSING COMPLETE!")
    print("="*50)

if __name__ == "__main__":
    main()