mod pose_detector;
mod visualization;
mod footstep_tracker;
mod person_detector;
mod udp;

use anyhow::{ Context, Result };
use clap::Parser;
use opencv::{ highgui, prelude::*, videoio };
use rayon::prelude::*;
use pose_detector::{ PoseDetector, PoseDetectorConfig, Keypoints };
use visualization::{ draw_footsteps, draw_bounding_boxes };
use footstep_tracker::FootstepTracker;
use person_detector::{ YoloDetector, YoloDetectorConfig, BoundingBox, PersonTracker };
use std::sync::{ Arc, Mutex };
use udp::UdpSender;

#[cfg(feature = "debug")]
use tracing::{ debug, error, info };
#[cfg(feature = "debug")]
use tracing_subscriber::EnvFilter;
#[cfg(feature = "debug")]
use indicatif::{ ProgressBar, ProgressStyle };

#[cfg(feature = "debug")]
fn init_tracing() {
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));

    // Ignore errors if a subscriber was already set by tests or embedding.
    let _ = tracing_subscriber::fmt().with_env_filter(filter).with_ansi(true).try_init();
}

#[derive(Clone, Debug)]
enum VideoSource {
    Webcam(i32),
    File(String),
}

#[derive(Parser, Debug)]
#[command(name = "footstep-tracker", about = "Footstep tracking", version, author)]
struct Args {
    /// Path to the CoreML pose detection model package
    #[arg(short, long, value_name = "MODEL_PATH")]
    model_path: Option<String>,

    /// Video source: camera ID (e.g., 0) or video file path
    #[arg(
        default_value = "0",
        value_name = "CAMERA_ID|VIDEO_FILE",
        value_parser = parse_video_source
    )]
    video: VideoSource,

    /// UDP target address for footstep events (host:port)
    #[arg(
        short,
        long,
        value_name = "UDP_TARGET",
        num_args = 0..=1,
        default_missing_value = "127.0.0.1:7000"
    )]
    udp_target: Option<String>,
}

fn parse_video_source(input: &str) -> Result<VideoSource, String> {
    if input.contains('.') || input.contains('/') {
        return Ok(VideoSource::File(input.to_string()));
    }

    match input.parse::<i32>() {
        Ok(id) => Ok(VideoSource::Webcam(id)),
        Err(_) => Ok(VideoSource::Webcam(0)),
    }
}

fn main() -> Result<()> {
    #[cfg(feature = "debug")]
    init_tracing();

    let args = Args::parse();
    let model_path = args.model_path.unwrap_or_else(|| "models/rtmpose.mlpackage".to_string());
    let video_source = args.video;

    // Configure and initialize the YOLO person detector
    let yolo_config = YoloDetectorConfig::default();
    let mut person_detector = YoloDetector::new(yolo_config).context(
        "Failed to initialize YOLO detector"
    )?;

    let udp_sender = if let Some(udp_target) = args.udp_target {
        let sender = UdpSender::new(&udp_target)?;
        #[cfg(feature = "debug")]
        info!("Footstep UDP target: {}", sender.target());
        Some(sender)
    } else {
        #[cfg(feature = "debug")]
        tracing::warn!("No UDP target specified; footstep events will not be sent over UDP");
        None
    };

    // Configure and initialize the pose detector
    let config = PoseDetectorConfig {
        model_path,
        ..Default::default()
    };

    let worker_count = std::thread
        ::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4);

    #[cfg(feature = "debug")]
    info!("Configuring {} pose detector workers", worker_count);

    #[cfg(feature = "debug")]
    let pb = ProgressBar::new(worker_count as u64).with_style(
        ProgressStyle::default_bar()
            .template("{prefix:.bold.blue} [{bar:40.cyan/blue}] {pos}/{len} {msg}")
            .unwrap()
            .progress_chars("=>-")
    );

    #[cfg(feature = "debug")]
    pb.set_prefix("Pose workers");

    let detector_pool: Arc<Mutex<Vec<PoseDetector>>> = Arc::new(
        Mutex::new(
            (0..worker_count)
                .map(|idx| {
                    let res = PoseDetector::new(config.clone());
                    #[cfg(feature = "debug")]
                    {
                        pb.inc(1);
                        if idx + 1 == worker_count {
                            pb.finish_with_message("ready");
                        }
                    }
                    res
                })
                .collect::<Result<Vec<_>>>()
                .context("Failed to initialize pose detector pool")?
        )
    );

    // Initialize footstep tracker (footsteps visible for 5 seconds)
    let mut footstep_tracker = FootstepTracker::new(1000);

    // Stable ID tracker for person boxes
    let mut id_tracker = PersonTracker::new(1000);

    // Open video source and get FPS for video files
    let (mut cap, video_fps) = match &video_source {
        VideoSource::Webcam(camera_id) => {
            #[cfg(feature = "debug")]
            info!("Opening camera {}...", camera_id);
            let cap = videoio::VideoCapture
                ::new(*camera_id, videoio::CAP_AVFOUNDATION)
                .context("Failed to open video capture")?;

            if !videoio::VideoCapture::is_opened(&cap)? {
                #[cfg(feature = "debug")]
                {
                    error!("Error: Could not open camera {}", camera_id);
                    error!(
                        "Try a different camera ID with: cargo run --release -- <model_path> <camera_id>"
                    );
                }
                return Ok(());
            }
            (cap, None)
        }
        VideoSource::File(path) => {
            #[cfg(feature = "debug")]
            info!("Opening video file: {}", path);
            let cap = videoio::VideoCapture
                ::from_file(path, videoio::CAP_ANY)
                .context("Failed to open video file")?;

            if !videoio::VideoCapture::is_opened(&cap)? {
                #[cfg(feature = "debug")]
                {
                    error!("Error: Could not open video file: {}", path);
                    error!("Make sure the file exists and is a valid video format");
                }
                return Ok(());
            }

            let fps = cap.get(videoio::CAP_PROP_FPS)?;
            let frame_count = cap.get(videoio::CAP_PROP_FRAME_COUNT)?;
            #[cfg(feature = "debug")]
            info!("Video: {:.1} FPS, {} frames", fps, frame_count as i32);
            (cap, Some(fps))
        }
    };

    let is_video_file = matches!(video_source, VideoSource::File(_));

    // Calculate delay for video playback (in milliseconds)
    let frame_delay = if let Some(fps) = video_fps {
        if fps > 0.0 {
            (1000.0 / fps) as i32
        } else {
            30 // fallback to ~30 FPS
        }
    } else {
        1 // For webcam, minimal delay
    };

    #[cfg(feature = "debug")]
    info!("Press 'q' to quit");

    let window_name = "Footstep Tracker - CoreML";
    highgui::named_window(window_name, highgui::WINDOW_AUTOSIZE)?;

    let mut frame_index: usize = 0;

    loop {
        let mut frame = Mat::default();
        cap.read(&mut frame)?;

        if frame.empty() {
            // If it's a video file, loop back to the beginning
            if is_video_file {
                #[cfg(feature = "debug")]
                debug!("Looping video...");
                cap.set(videoio::CAP_PROP_POS_FRAMES, 0.0)?;
                cap.read(&mut frame)?;

                // If still empty after reset, break
                if frame.empty() {
                    break;
                }
            } else {
                break;
            }
        }

        // Stage 1: Detect people with YOLO
        #[cfg(feature = "debug")]
        let start = std::time::Instant::now();

        let people = person_detector.detect_people(&frame)?;
        let people_with_ids = id_tracker.assign_ids(people);

        #[cfg(feature = "debug")]
        {
            let yolo_time = start.elapsed();
            if frame_index % 30 == 0 {
                debug!("YOLO detection time: {:.2?}", yolo_time);
            }
        }

        // Stage 2: Run pose detection on each person's bounding box
        let pose_config = config.clone();
        let detector_pool = detector_pool.clone();

        let processed: Vec<(BoundingBox, Option<(usize, Keypoints)>)> = people_with_ids
            .par_iter()
            .enumerate()
            .map(
                |(person_idx, (person_id, bbox))| -> Result<_> {
                    let expanded_bbox = bbox.expand(1.4);

                    let person_crop = YoloDetector::crop_region(
                        &frame,
                        &expanded_bbox,
                        pose_config.input_height
                    )?;

                    let mut detector = match (
                        {
                            let mut pool = detector_pool.lock().unwrap();
                            pool.pop()
                        }
                    ) {
                        Some(det) => det,
                        None =>
                            PoseDetector::new(pose_config.clone()).context(
                                "Failed to initialize pose detector for worker"
                            )?,
                    };

                    let mut keyed = None;
                    if let Ok(mut person_keypoints_batch) = detector.detect_pose(&person_crop) {
                        if let Some(person_keypoints) = person_keypoints_batch.first_mut() {
                            let (x, y, w, h) = expanded_bbox.to_pixels(frame.cols(), frame.rows());

                            for kp in person_keypoints.iter_mut() {
                                kp[1] = (kp[1] * (w as f32) + (x as f32)) / (frame.cols() as f32);
                                kp[0] = (kp[0] * (h as f32) + (y as f32)) / (frame.rows() as f32);
                            }

                            keyed = Some((*person_id, person_keypoints.clone()));
                        }

                        #[cfg(feature = "debug")]
                        {
                            if frame_index % 30 == 0 {
                                debug!(
                                    "Person {}: bbox confidence={:.2}",
                                    person_idx,
                                    bbox.confidence
                                );
                            }
                        }
                    }

                    {
                        let mut pool = detector_pool.lock().unwrap();
                        pool.push(detector);
                    }

                    Ok((expanded_bbox, keyed))
                }
            )
            .collect::<Result<Vec<_>>>()?;

        let expanded_bboxes: Vec<BoundingBox> = processed
            .iter()
            .map(|(bbox, _)| bbox.clone())
            .collect();

        let keyed_keypoints: Vec<(usize, Keypoints)> = processed
            .into_iter()
            .filter_map(|(_, keyed)| keyed)
            .collect();

        #[cfg(feature = "debug")]
        {
            if frame_index % 30 == 0 {
                let total_time = start.elapsed();
                debug!(
                    "Total detection time (YOLO + {} poses): {:.2?}",
                    people_with_ids.len(),
                    total_time
                );
            }
        }

        // Update footstep tracking
        let new_footsteps = footstep_tracker.update(&keyed_keypoints);

        if let Some(sender) = udp_sender.as_ref() {
            for event in &new_footsteps {
                #[cfg(feature = "debug")]
                if let Err(err) = sender.send(event) {
                    error!("Failed to send UDP footstep packet: {:?}", err);
                }

                #[cfg(not(feature = "debug"))]
                {
                    let _ = sender.send(event);
                }
            }
        }

        // Draw expanded bounding boxes (used for pose detection)
        draw_bounding_boxes(&mut frame, &people_with_ids)?;

        // Visualize all keypoints
        // draw_all_keypoints(&mut frame, &all_keypoints, 0.1)?;
        // draw_all_ankles(&mut frame, &all_keypoints, 0.1)?;

        // Draw footsteps
        let active_footsteps = footstep_tracker.get_all_footsteps();
        let archived = footstep_tracker.get_archived_footsteps();
        draw_footsteps(&mut frame, &active_footsteps, archived)?;

        highgui::imshow(window_name, &frame)?;

        // Break on 'q' key - use appropriate delay for video playback
        if highgui::wait_key(1)? == (b'q' as i32) {
            break;
        }

        frame_index = frame_index.wrapping_add(1);
    }

    Ok(())
}
