//! Image Classifier Node
//!
//! Hardware-agnostic image classification node that:
//! - Subscribes to camera topics (simulation/real/fake - doesn't matter!)
//! - Loads ONNX models from catalog with preprocessing metadata
//! - Runs inference and publishes classification results
//!
//! # Topic Interface
//!
//! **Input:** `/robot/sensors/camera/rgb` (CameraImage)
//! **Output:** `/vision/classification` (Classification)
//!
//! # Hardware Agnosticism
//!
//! This node doesn't care about the camera source:
//! - Simulation → simulation-bridge publishes to Redis
//! - Real camera → camera driver publishes to Redis
//! - Fake data → fake-camera publishes to Redis
//!
//! Same topic = same node works everywhere!

mod config;

pub use config::{ImageClassifierConfig, ModelCatalogEntry, ModelConfig};

use anyhow::{Context as AnyhowContext, Result};
use base64::Engine;
use image::{DynamicImage, ImageFormat};
use mecha10_core::messages::Message;
use mecha10_core::model::{load_labels, load_model_config, ModelConfig as CoreModelConfig};
use mecha10_core::prelude::*;
use mecha10_core::topics::Topic;
use ort::session::Session;
use serde::{Deserialize, Serialize};
use std::io::Cursor;
use std::path::PathBuf;
use tracing::info;

// === Message Types ===

/// Camera image message (hardware agnostic!)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CameraImage {
    pub camera_id: String,
    pub width: u32,
    pub height: u32,
    pub timestamp: u64,
    /// Raw image bytes (JPEG or RGB)
    #[serde(with = "serde_bytes")]
    pub image_bytes: Vec<u8>,
    /// Image format ("jpeg" or "rgb")
    pub format: String,
}

impl Message for CameraImage {}

/// Classification result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Classification {
    pub class: String,
    pub class_id: usize,
    pub confidence: f32,
    pub top_k: Vec<ClassScore>,
    pub timestamp: u64,
    /// Base64-encoded JPEG thumbnail of the classified image (optional)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub image_base64: Option<String>,
}

impl Message for Classification {}

/// Individual class score
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassScore {
    pub class: String,
    pub class_id: usize,
    pub confidence: f32,
}

/// Classification control commands
///
/// Control topic: /vision/classification/control
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "command", rename_all = "snake_case")]
pub enum ClassificationControl {
    /// Start continuous classification
    Start {
        #[serde(default)]
        rate_hz: Option<u32>,
    },
    /// Stop continuous classification
    Stop,
    /// Classify a single frame (one-shot)
    SingleShot,
}

impl Message for ClassificationControl {}

/// Classification mode
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ClassificationMode {
    /// Idle - not classifying
    Idle,
    /// Continuous classification at configured rate
    Continuous,
    /// Single-shot requested - classify next frame only
    SingleShot,
}

// === Image Classifier Node ===

pub struct ImageClassifierNode {
    config: ImageClassifierConfig,
    session: Option<Session>,
    frame_count: u64,
    mode: ClassificationMode,
    /// Class labels loaded from model directory
    labels: Vec<String>,
    /// Model configuration loaded from model directory
    model_config: CoreModelConfig,
}

impl ImageClassifierNode {
    pub fn new(config: ImageClassifierConfig, labels: Vec<String>, model_config: CoreModelConfig) -> Self {
        Self {
            config,
            session: None,
            frame_count: 0,
            mode: ClassificationMode::Idle, // Start idle, wait for trigger
            labels,
            model_config,
        }
    }

    /// Handle classification control command
    fn handle_control(&mut self, control: ClassificationControl) {
        tracing::info!("🎮 Received control command: {:?}", control);
        match control {
            ClassificationControl::Start { rate_hz } => {
                self.mode = ClassificationMode::Continuous;
                if let Some(rate) = rate_hz {
                    self.config.processing_rate_hz = rate;
                }
                tracing::info!(
                    "✅ Started continuous classification at {} Hz",
                    self.config.processing_rate_hz
                );
            }
            ClassificationControl::Stop => {
                self.mode = ClassificationMode::Idle;
                tracing::info!("⏸️  Stopped classification");
            }
            ClassificationControl::SingleShot => {
                self.mode = ClassificationMode::SingleShot;
                tracing::info!("📸 Single-shot classification requested");
            }
        }
    }

    /// Check if we should process this frame based on mode and rate limiting
    fn should_process_frame(&mut self) -> bool {
        tracing::info!("🔍 should_process_frame called, current mode: {:?}", self.mode);
        let result = match self.mode {
            ClassificationMode::Idle => {
                tracing::info!("   Mode is Idle, returning false");
                false
            }
            ClassificationMode::SingleShot => {
                tracing::info!("   Mode is SingleShot, switching to Idle and returning true");
                // Process once, then go back to idle
                self.mode = ClassificationMode::Idle;
                true
            }
            ClassificationMode::Continuous => {
                self.frame_count += 1;
                let skip_frames = 30 / self.config.processing_rate_hz.max(1);
                let should_process = self.frame_count % skip_frames as u64 == 0;
                tracing::info!(
                    "   Mode is Continuous, frame_count={}, skip_frames={}, should_process={}",
                    self.frame_count,
                    skip_frames,
                    should_process
                );
                should_process
            }
        };
        tracing::info!("   Returning: {}", result);
        result
    }

    /// Process a camera image and return classification
    async fn classify_image(&mut self, msg: &CameraImage) -> Result<Classification> {
        let session = self.session.as_mut().context("ONNX session not initialized")?;

        // 1. Decode image
        let image = decode_camera_image(msg)?;

        // 2. Create thumbnail for visualization (keep original aspect ratio)
        let thumbnail = image.resize(320, 240, image::imageops::FilterType::Lanczos3);
        let image_base64 = encode_image_as_base64_jpeg(&thumbnail)?;

        // 3. Preprocess (resize + normalize) using model config
        let [width, height] = self.model_config.input_size;
        let input_tensor = preprocess_image(
            image,
            width,
            height,
            &self.model_config.preprocessing.mean,
            &self.model_config.preprocessing.std,
        )?;

        // 4. Run ONNX inference
        // Create tensor value from array view
        let input_value = ort::value::TensorRef::from_array_view(input_tensor.view())?;

        // Extract logits in a separate scope to drop outputs and release the borrow
        let probabilities = {
            let outputs = session.run(ort::inputs![input_value])?;

            // 5. Extract output tensor (shape: [1, 1000] for ImageNet classification)
            let logits_array: ndarray::ArrayViewD<f32> = outputs[0]
                .try_extract_array()
                .context("Failed to extract output tensor as array")?;
            let logits_vec: Vec<f32> = logits_array
                .as_slice()
                .context("Failed to get tensor data as slice")?
                .to_vec();

            // 6. Apply softmax to get probabilities
            softmax(&logits_vec)
        }; // outputs dropped here, releasing the mutable borrow on session

        // 7. Get top-K predictions (now we can borrow self immutably)
        let top_k = self.top_k_predictions(&probabilities, self.config.top_k);

        // 8. Create classification result
        let best = &top_k[0];
        Ok(Classification {
            class: best.class.clone(),
            class_id: best.class_id,
            confidence: best.confidence,
            top_k,
            timestamp: msg.timestamp,
            image_base64: Some(image_base64),
        })
    }

    /// Get top K predictions with class names
    fn top_k_predictions(&self, probabilities: &[f32], k: usize) -> Vec<ClassScore> {
        let mut indexed: Vec<(usize, f32)> = probabilities.iter().enumerate().map(|(i, &p)| (i, p)).collect();

        // Sort by probability (descending)
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        // Take top K
        indexed
            .into_iter()
            .take(k)
            .map(|(idx, prob)| ClassScore {
                class: self
                    .labels
                    .get(idx)
                    .cloned()
                    .unwrap_or_else(|| format!("unknown_class_{}", idx)),
                class_id: idx,
                confidence: prob,
            })
            .collect()
    }
}

// === Utility Functions ===

/// Decode camera image from various formats
fn decode_camera_image(msg: &CameraImage) -> Result<DynamicImage> {
    // Use raw image bytes directly
    let bytes = &msg.image_bytes;

    // Decode image format
    let image = match msg.format.as_str() {
        "jpeg" | "jpg" => image::load_from_memory_with_format(bytes, ImageFormat::Jpeg)?,
        "png" => image::load_from_memory_with_format(bytes, ImageFormat::Png)?,
        "rgb" | "rgb8" => {
            let img = image::RgbImage::from_raw(msg.width, msg.height, bytes.clone())
                .context("Failed to create RGB image from raw bytes")?;
            DynamicImage::ImageRgb8(img)
        }
        "rgba" | "rgba8" => {
            let img = image::RgbaImage::from_raw(msg.width, msg.height, bytes.clone())
                .context("Failed to create RGBA image from raw bytes")?;
            DynamicImage::ImageRgba8(img)
        }
        _ => anyhow::bail!("Unsupported image format: {}", msg.format),
    };

    Ok(image)
}

/// Encode image as base64-encoded JPEG for transmission to dashboard
fn encode_image_as_base64_jpeg(image: &DynamicImage) -> Result<String> {
    // Encode image as JPEG with quality 85
    let mut buffer = Cursor::new(Vec::new());
    image
        .write_to(&mut buffer, ImageFormat::Jpeg)
        .context("Failed to encode image as JPEG")?;

    // Convert to base64
    let base64_string = base64::engine::general_purpose::STANDARD.encode(buffer.into_inner());

    Ok(base64_string)
}

/// Preprocess image for model inference with dynamic normalization
fn preprocess_image(
    image: DynamicImage,
    width: u32,
    height: u32,
    mean: &[f32; 3],
    std: &[f32; 3],
) -> Result<ndarray::Array4<f32>> {
    // Resize to model input size
    let resized = image.resize_exact(width, height, image::imageops::FilterType::Lanczos3);

    // Convert to RGB
    let rgb = resized.to_rgb8();

    // Create tensor [1, 3, H, W]
    let mut array = ndarray::Array4::<f32>::zeros((1, 3, height as usize, width as usize));

    for (x, y, pixel) in rgb.enumerate_pixels() {
        let r = (pixel[0] as f32 / 255.0 - mean[0]) / std[0];
        let g = (pixel[1] as f32 / 255.0 - mean[1]) / std[1];
        let b = (pixel[2] as f32 / 255.0 - mean[2]) / std[2];

        array[[0, 0, y as usize, x as usize]] = r;
        array[[0, 1, y as usize, x as usize]] = g;
        array[[0, 2, y as usize, x as usize]] = b;
    }

    Ok(array)
}

/// Apply softmax to convert logits to probabilities
fn softmax(logits: &[f32]) -> Vec<f32> {
    let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exp: Vec<f32> = logits.iter().map(|&x| (x - max).exp()).collect();
    let sum: f32 = exp.iter().sum();
    exp.iter().map(|&x| x / sum).collect()
}

/// Load model from structured directory or path
fn load_model(config: &ModelConfig) -> Result<PathBuf> {
    if let Some(name) = &config.name {
        // Use catalog model from structured directory
        let path = PathBuf::from("models").join(name).join("model.onnx");
        if !path.exists() {
            anyhow::bail!(
                "Model '{}' not found at {}. Run: mecha10 models pull {}",
                name,
                path.display(),
                name
            );
        }
        Ok(path)
    } else if let Some(path) = &config.path {
        // Use direct path
        if !path.exists() {
            anyhow::bail!("Model file not found: {}", path.display());
        }
        Ok(path.clone())
    } else {
        anyhow::bail!("Either model.name or model.path must be specified");
    }
}

// === Node Entry Point ===

pub async fn run() -> Result<()> {
    // Note: Tracing is initialized by robot-dev launcher, not here
    info!("🤖 Starting Image Classifier Node");

    // Create context first (needed for config loading)
    let ctx = Context::new("image-classifier").await?;

    // Load config using V2 system with environment-aware file-level fallback
    let config: ImageClassifierConfig = ctx.load_node_config("image-classifier").await?;
    info!("Configuration: {:?}", config);

    // Load class labels and preprocessing params from model directory
    let model_name = config
        .model
        .name
        .as_ref()
        .context("model.name is required. Custom model paths must follow the structured format: models/<name>/{model.onnx, labels.txt, config.json}")?;

    info!("📋 Loading model configuration for: {}", model_name);

    // Use shared functions from mecha10-core
    let labels = load_labels(model_name)?;
    let model_config = load_model_config(model_name)?;

    info!("✅ Loaded {} class labels", labels.len());
    info!(
        "✅ Loaded preprocessing: input_size={:?}, mean={:?}, std={:?}",
        model_config.input_size, model_config.preprocessing.mean, model_config.preprocessing.std
    );

    // Load ONNX model
    let model_path = load_model(&config.model)?;
    info!("Loading model from: {}", model_path.display());

    let session = Session::builder()?.commit_from_file(&model_path)?;

    info!("✅ Model loaded successfully");

    // Subscribe to camera topic
    let input_topic = config.input_topic();
    let output_topic = config.output_topic();
    let control_topic_path = config.control_topic();

    info!("📡 Subscribing to: {}", input_topic);
    info!("📡 Control topic: {}", control_topic_path);
    info!("📤 Publishing to: {}", output_topic);
    info!("⚡ Default processing rate: {} Hz", config.processing_rate_hz);
    info!("⏸️  Starting in IDLE mode - waiting for trigger");

    // Create node state
    let mut node = ImageClassifierNode::new(config.clone(), labels, model_config);
    node.session = Some(session);

    // Create topics (using Box::leak for dynamic topic paths)
    // Note: This is acceptable for topic paths that are loaded once at startup
    let input_topic_static: &'static str = Box::leak(input_topic.into_boxed_str());
    let output_topic_static: &'static str = Box::leak(output_topic.into_boxed_str());
    let control_topic_static: &'static str = Box::leak(control_topic_path.into_boxed_str());

    let control_topic = Topic::<ClassificationControl>::new(control_topic_static);
    // Camera frames come wrapped in RedisMessage envelope from simulation-bridge
    let camera_topic = Topic::<mecha10_core::messages::RedisMessage<CameraImage>>::new(input_topic_static);
    let output_topic = Topic::<Classification>::new(output_topic_static);

    // Subscribe to topics
    let mut control_receiver = ctx.subscribe(control_topic).await?;
    let mut camera_receiver = ctx.subscribe(camera_topic).await?;

    info!("✅ Subscribed to control and camera topics");

    // Main processing loop
    info!("🔄 Entering main processing loop");
    loop {
        tokio::select! {
            // Handle control commands
            Some(control_msg) = control_receiver.recv() => {
                info!("📨 Got control message from Redis");
                node.handle_control(control_msg);
            }

            // Handle camera frames
            Some(camera_envelope) = camera_receiver.recv() => {
                // Extract CameraImage from RedisMessage envelope
                let camera_msg = camera_envelope.payload;

                info!("📷 Got camera frame: {}x{} format={} (mode={:?})",
                    camera_msg.width, camera_msg.height, camera_msg.format, node.mode);

                let should_process = node.should_process_frame();
                info!("🤔 should_process_frame() returned: {} (mode={:?})", should_process, node.mode);

                if should_process {
                    info!("🔍 Processing frame for classification");
                    match node.classify_image(&camera_msg).await {
                        Ok(classification) => {
                            info!("✅ {} ({:.1}%)",
                                classification.class,
                                classification.confidence * 100.0
                            );

                            // Publish classification result
                            ctx.publish_to(output_topic, &classification).await?;
                            info!("📤 Published classification result to {}", output_topic.path());
                        }
                        Err(e) => {
                            tracing::warn!("❌ Classification failed: {}", e);
                        }
                    }
                } else {
                    info!("⏭️  Skipping frame (mode={:?})", node.mode);
                }
            }
        }
    }
}
