use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

use crate::sfu::ForwardPacket;

/// AI Stream Processor for running STT, TTS, and Vision models directly on the edge.
/// By processing media at the relay node through localized ML services, we
/// avoid the latency of sending packets back to a centralized cloud AI service.
pub struct AiProcessor {
    pub config: Arc<RwLock<AiConfig>>,
    pub client: reqwest::Client,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AiConfig {
    pub stt_endpoint: Option<String>,
    pub tts_endpoint: Option<String>,
    pub vision_endpoint: Option<String>,
}

#[derive(Deserialize)]
struct SttResponse {
    text: String,
}

#[derive(Deserialize)]
struct VisionResponse {
    objects: Vec<BoundingBox>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct BoundingBox {
    pub label: String,
    pub confidence: f32,
    pub x: u32,
    pub y: u32,
    pub w: u32,
    pub h: u32,
}

impl AiProcessor {
    pub fn new(config: AiConfig) -> Self {
        info!("[AI Processor] Initializing Edge AI Inference Client...");

        Self {
            config: Arc::new(RwLock::new(config)),
            client: reqwest::Client::new(),
        }
    }

    /// Pre-load necessary AI models or establish persistent inference server connections.
    pub async fn load_models(&self) -> Result<(), &'static str> {
        let conf = self.config.read().await;

        if let Some(endpoint) = &conf.stt_endpoint {
            info!(
                "[AI Processor] Checking STT availability at {}...",
                endpoint
            );
            // Example: ping or warm-up the STT worker
        } else {
            debug!("[AI Processor] STT endpoint not configured. STT disabled.");
        }

        if let Some(endpoint) = &conf.vision_endpoint {
            info!(
                "[AI Processor] Checking Vision model availability at {}...",
                endpoint
            );
        } else {
            debug!("[AI Processor] Vision endpoint not configured. Vision disabled.");
        }

        Ok(())
    }

    /// Process a live streaming packet and apply AI inference (STT).
    pub async fn process_audio_stt(&self, packet: &ForwardPacket) -> Option<String> {
        let conf = self.config.read().await;
        let endpoint = conf.stt_endpoint.as_ref()?;

        debug!(
            "[AI Processor] Running STT on packet of size {} bytes via {}...",
            packet.payload.len(),
            endpoint
        );

        let req = self
            .client
            .post(endpoint)
            .header("Content-Type", "application/octet-stream")
            .body(packet.payload.clone());

        match req.send().await {
            Ok(resp) => {
                if let Ok(json) = resp.json::<SttResponse>().await {
                    Some(json.text)
                } else {
                    None
                }
            }
            Err(e) => {
                warn!("[AI Processor] STT request failed: {}", e);
                None
            }
        }
    }

    /// Process a live streaming packet and apply Computer Vision.
    pub async fn process_video_vision(&self, packet: &ForwardPacket) -> Option<Vec<BoundingBox>> {
        let conf = self.config.read().await;
        let endpoint = conf.vision_endpoint.as_ref()?;

        debug!(
            "[AI Processor] Running Object Detection on packet sequence {} via {}...",
            packet.sequence, endpoint
        );

        let req = self
            .client
            .post(endpoint)
            .header("Content-Type", "application/octet-stream")
            .body(packet.payload.clone());

        match req.send().await {
            Ok(resp) => {
                if let Ok(json) = resp.json::<VisionResponse>().await {
                    Some(json.objects)
                } else {
                    None
                }
            }
            Err(e) => {
                warn!("[AI Processor] Vision request failed: {}", e);
                None
            }
        }
    }
}
