//! Native GUI for the 360° -> COLMAP cubemap rig pipeline.
//!
//! Built with **egui** (rendered through **wgpu** via eframe) so it is
//! cross-vendor / cross-platform (Vulkan/Metal/DX12) just like the pipeline.
//! The GUI drives the pipeline by shelling out to the `pipeline_360` CLI binary
//! with `--gui-log-output`, streaming its `[progress] ` / log line protocol —
//! the same protocol the Python GUI uses — so the heavy GPU work runs in the
//! (optionally release-built) CLI process and the GUI stays responsive.
//!
//! Build:  cargo build --features gui
//! Run:    ./target/debug/pipeline_360_gui
//!
//! The GUI looks for the `pipeline_360` binary beside itself, then on PATH.

//! Native GUI (egui + wgpu). Enabled by the `gui` feature.
#![cfg(feature = "gui")]

use std::collections::VecDeque;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::mpsc::{self, Receiver, Sender};
use std::time::Instant;

use eframe::egui;

const PROGRESS_PREFIX: &str = "[progress] ";
const LOG_CAP: usize = 4000;

/// Locate the `pipeline_360` CLI binary: beside the GUI exe, else on PATH.
fn resolve_cli_binary() -> Option<PathBuf> {
    if let Ok(exe) = std::env::current_exe() {
        if let Some(dir) = exe.parent() {
            for name in ["pipeline_360", "pipeline_360.exe"] {
                let p = dir.join(name);
                if p.exists() {
                    return Some(p);
                }
            }
        }
    }
    which("pipeline_360")
}

/// Minimal `which` (no extra dep): search PATH for an executable.
fn which(name: &str) -> Option<PathBuf> {
    let path = std::env::var_os("PATH")?;
    for dir in std::env::split_paths(&path) {
        let full = dir.join(name);
        if full.is_file() {
            return Some(full);
        }
    }
    None
}

enum Line {
    Log(String),
    Progress(String),
    Done(i32),
}

struct RunHandle {
    child: std::process::Child,
    rx: Receiver<Line>,
    started: Instant,
    code: Option<i32>,
}

impl RunHandle {
    fn spawn(mut cmd: std::process::Command) -> std::io::Result<RunHandle> {
        let mut child = cmd
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;
        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();
        let (tx, rx) = mpsc::channel::<Line>();
        spawn_line_reader(stdout, tx.clone(), false);
        spawn_line_reader(stderr, tx, true);
        Ok(RunHandle { child, rx, started: Instant::now(), code: None })
    }

    fn poll(&mut self, log: &mut VecDeque<String>) {
        while let Ok(line) = self.rx.try_recv() {
            match line {
                Line::Log(s) | Line::Progress(s) => {
                    log.push_back(s);
                    while log.len() > LOG_CAP {
                        log.pop_front();
                    }
                }
                Line::Done(c) => self.code = Some(c),
            }
        }
        if self.code.is_none() {
            if let Ok(Some(status)) = self.child.try_wait() {
                self.code = Some(status.code().unwrap_or(-1));
            }
        }
    }

    fn kill(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
        self.code = Some(-1);
    }

    fn elapsed(&self) -> std::time::Duration {
        self.started.elapsed()
    }
}

fn spawn_line_reader<R: std::io::Read + Send + 'static>(mut r: R, tx: Sender<Line>, is_stderr: bool) {
    std::thread::spawn(move || {
        let mut buf = [0u8; 8192];
        let mut leftover = Vec::new();
        loop {
            let n = match r.read(&mut buf) {
                Ok(0) => break,
                Ok(n) => n,
                Err(_) => break,
            };
            leftover.extend_from_slice(&buf[..n]);
            while let Some(pos) = leftover.iter().position(|&b| b == b'\n') {
                let line_bytes = leftover.drain(..=pos).collect::<Vec<u8>>();
                let raw = String::from_utf8_lossy(&line_bytes);
                let raw = raw.trim_end_matches(['\n', '\r']);
                if let Some(rest) = raw.strip_prefix(PROGRESS_PREFIX) {
                    let _ = tx.send(Line::Progress(format!("[…] {rest}")));
                } else if !raw.is_empty() {
                    let _ = tx.send(Line::Log(if is_stderr {
                        format!("[stderr] {raw}")
                    } else {
                        raw.to_string()
                    }));
                }
            }
        }
        let _ = tx.send(Line::Done(0));
    });
}

// ---------------------------------------------------------------------------
// Form state — mirrors the CLI flags + defaults.
// ---------------------------------------------------------------------------

#[derive(Clone)]
struct ExtractForm {
    chunk_size: u32,
    scale_width: u32,
    output_format: String,
    jpeg_quality: u32,
    min_sharpness: f64,
    min_exposure: f64,
    max_exposure: f64,
    min_contrast: f64,
    min_motion: f64,
    max_motion: f64,
    selection: String,
    target_count: u32,
    min_gap: u32,
    bad_frame_action: String,
    motion_mode: String,
    batch: u32,
}

impl Default for ExtractForm {
    fn default() -> Self {
        Self {
            chunk_size: 30, scale_width: 1920, output_format: "png".into(), jpeg_quality: 95,
            min_sharpness: 0.0, min_exposure: 0.0, max_exposure: 255.0, min_contrast: 0.0,
            min_motion: 0.0, max_motion: 1.0e12, selection: "chunk".into(), target_count: 0,
            min_gap: 0, bad_frame_action: "keep".into(), motion_mode: "frame_diff".into(), batch: 8,
        }
    }
}

#[derive(Clone)]
struct MasksForm {
    classes: String,
    model: PathBuf,
    device: String,
    conf: f64,
    imgsz: u32,
    dilate: u32,
    projection: String,
    face_size: u32,
    skip_faces: String,
    tile_batch: u32,
    empty_if_none: bool,
    iou: f64,
}

impl Default for MasksForm {
    fn default() -> Self {
        Self {
            classes: "person,car,bus,truck,bicycle,motorcycle".into(),
            model: PathBuf::from("yolo11x-seg.onnx"), device: "cpu".into(), conf: 0.35,
            imgsz: 1280, dilate: 6, projection: "cubemap".into(), face_size: 1024,
            skip_faces: String::new(), tile_batch: 8, empty_if_none: false, iou: 0.45,
        }
    }
}

#[derive(Clone)]
struct CubemapForm { face_size: u32, skip_faces: String, jpeg_quality: u32, batch: u32 }
impl Default for CubemapForm {
    fn default() -> Self { Self { face_size: 1600, skip_faces: "U,D".into(), jpeg_quality: 95, batch: 4 } }
}

#[derive(Clone)]
struct RigForm { ref_face: String, intrinsics_convention: String, link_mode: String }
impl Default for RigForm {
    fn default() -> Self { Self { ref_face: "F".into(), intrinsics_convention: "linspace".into(), link_mode: "symlink".into() } }
}

// ---------------------------------------------------------------------------
// App
// ---------------------------------------------------------------------------

struct App {
    project_dir: PathBuf,
    videos: Vec<PathBuf>,
    start_stage: String,
    skip_masks: bool,
    extract: ExtractForm,
    masks: MasksForm,
    cubemap: CubemapForm,
    rig: RigForm,
    cli_binary: Option<PathBuf>,
    cli_binary_str: String,
    run: Option<RunHandle>,
    log: VecDeque<String>,
    status: String,
    last_progress: String,
}

impl Default for App {
    fn default() -> Self {
        let cli = resolve_cli_binary();
        Self {
            project_dir: PathBuf::new(), videos: Vec::new(), start_stage: "extract".into(),
            skip_masks: false, extract: ExtractForm::default(), masks: MasksForm::default(),
            cubemap: CubemapForm::default(), rig: RigForm::default(),
            cli_binary_str: cli.as_ref().map(|p| p.display().to_string()).unwrap_or_else(|| "pipeline_360 (NOT FOUND — build it first)".into()),
            cli_binary: cli, log: VecDeque::new(), run: None, status: "Idle".into(), last_progress: String::new(),
        }
    }
}

impl App {
    fn is_running(&self) -> bool {
        self.run.as_ref().is_some_and(|h| h.code.is_none())
    }

    fn build_command(&self) -> Option<std::process::Command> {
        let bin = self.cli_binary.clone()?;
        let mut cmd = std::process::Command::new(bin);
        cmd.arg("all");
        cmd.arg("--project-dir").arg(&self.project_dir);
        for v in &self.videos {
            cmd.arg("--video").arg(v);
        }
        cmd.arg("--start-stage").arg(&self.start_stage);
        if self.skip_masks {
            cmd.arg("--skip-masks");
        }
        cmd.arg("--gui-log-output");

        let e = &self.extract;
        cmd.args([
            "--extract-chunk-size", &e.chunk_size.to_string(),
            "--extract-scale-width", &e.scale_width.to_string(),
            "--extract-output-format", &e.output_format,
            "--extract-jpeg-quality", &e.jpeg_quality.to_string(),
            "--extract-min-sharpness", &fmt_f(e.min_sharpness),
            "--extract-min-exposure", &fmt_f(e.min_exposure),
            "--extract-max-exposure", &fmt_f(e.max_exposure),
            "--extract-min-contrast", &fmt_f(e.min_contrast),
            "--extract-min-motion", &fmt_f(e.min_motion),
            "--extract-max-motion", &fmt_f(e.max_motion),
            "--extract-selection", &e.selection,
            "--extract-target-count", &e.target_count.to_string(),
            "--extract-min-gap", &e.min_gap.to_string(),
            "--extract-bad-frame-action", &e.bad_frame_action,
            "--extract-motion-mode", &e.motion_mode,
            "--extract-batch", &e.batch.to_string(),
        ]);

        let m = &self.masks;
        cmd.args([
            "--mask-classes", &m.classes,
            "--mask-model", &m.model.display().to_string(),
            "--mask-device", &m.device,
            "--mask-conf", &fmt_f(m.conf),
            "--mask-imgsz", &m.imgsz.to_string(),
            "--mask-dilate", &m.dilate.to_string(),
            "--mask-projection", &m.projection,
            "--mask-face-size", &m.face_size.to_string(),
            "--mask-skip-faces", &m.skip_faces,
            "--mask-tile-batch", &m.tile_batch.to_string(),
            "--mask-iou", &fmt_f(m.iou),
        ]);
        if m.empty_if_none {
            cmd.arg("--mask-empty-if-none");
        }

        let c = &self.cubemap;
        cmd.args([
            "--cubemap-face-size", &c.face_size.to_string(),
            "--cubemap-skip-faces", &c.skip_faces,
            "--cubemap-jpeg-quality", &c.jpeg_quality.to_string(),
            "--cubemap-batch", &c.batch.to_string(),
        ]);

        let r = &self.rig;
        cmd.args([
            "--rig-ref-face", &r.ref_face,
            "--rig-intrinsics-convention", &r.intrinsics_convention,
            "--rig-link-mode", &r.link_mode,
        ]);
        Some(cmd)
    }

    fn start(&mut self) {
        if self.is_running() {
            return;
        }
        if self.cli_binary.is_none() {
            self.status = "Error: pipeline_360 binary not found (build it first).".into();
            return;
        }
        if self.project_dir.as_os_str().is_empty() || self.videos.is_empty() {
            self.status = "Error: choose a project dir and at least one video.".into();
            return;
        }
        let Some(cmd) = self.build_command() else { return };
        self.log.clear();
        self.last_progress.clear();
        self.status = "Running…".into();
        self.log.push_back(format!("$ {}", cmd_to_string(&cmd)));
        match RunHandle::spawn(cmd) {
            Ok(h) => self.run = Some(h),
            Err(e) => self.status = format!("Error: failed to start pipeline: {e}"),
        }
    }

    fn cancel(&mut self) {
        if let Some(h) = self.run.as_mut() {
            h.kill();
            self.status = "Cancelled.".into();
        }
    }
}

fn fmt_f(v: impl ToString) -> String {
    let s = v.to_string();
    if s.eq_ignore_ascii_case("inf") { "1e12".into() } else { s }
}

fn cmd_to_string(cmd: &std::process::Command) -> String {
    let prog = cmd.get_program().to_string_lossy().to_string();
    let args: Vec<String> = cmd.get_args().map(|a| {
        let a = a.to_string_lossy().to_string();
        if a.contains(' ') { format!("\"{a}\"") } else { a }
    }).collect();
    format!("{prog} {}", args.join(" "))
}

fn combo(ui: &mut egui::Ui, id: &str, options: &[&str], val: &mut String) {
    egui::ComboBox::from_id_salt(id).selected_text(val.clone()).show_ui(ui, |ui| {
        for o in options {
            ui.selectable_value(val, (*o).to_string(), *o);
        }
    });
}

impl eframe::App for App {
    fn ui(&mut self, ui: &mut egui::Ui, _frame: &mut eframe::Frame) {
        // Poll subprocess each frame.
        if let Some(h) = self.run.as_mut() {
            h.poll(&mut self.log);
            if let Some(code) = h.code {
                if !self.status.starts_with("Cancelled") {
                    self.status = if code == 0 { "Pipeline complete.".into() } else { format!("Pipeline exited with code {code}. See log.") };
                }
            }
        }

        // Bottom: log panel + run controls.
        egui::Panel::bottom("log_panel")
            .resizable(true)
            .min_size(150.0)
            .show(ui, |ui| {
                ui.horizontal(|ui| {
                    let status = if self.is_running() {
                        if let Some(h) = self.run.as_ref() {
                            format!("{} ({}s)", self.status, h.elapsed().as_secs())
                        } else { self.status.clone() }
                    } else { self.status.clone() };
                    ui.label(egui::RichText::new(status).strong());
                    ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                        let running = self.is_running();
                        if running {
                            ui.spinner();
                        }
                        if ui.add_enabled(running, egui::Button::new("⏹ Stop")).clicked() {
                            self.cancel();
                        }
                        if ui.add_enabled(!running, egui::Button::new("▶ Run pipeline")).clicked() {
                            self.start();
                        }
                    });
                });
                ui.separator();
                egui::ScrollArea::vertical()
                    .auto_shrink([false, false])
                    .stick_to_bottom(true)
                    .show(ui, |ui| {
                        ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Wrap);
                        for line in &self.log {
                            if line.starts_with("[…] ") {
                                continue;
                            }
                            ui.label(egui::RichText::new(line).monospace().size(12.0));
                        }
                        if !self.last_progress.is_empty() {
                            ui.label(
                                egui::RichText::new(&self.last_progress)
                                    .monospace()
                                    .color(egui::Color32::LIGHT_BLUE)
                                    .size(12.0),
                            );
                        }
                    });
            });

        // Left: inputs.
        egui::Panel::left("inputs")
            .resizable(true)
            .min_size(320.0)
            .show(ui, |ui| {
                egui::ScrollArea::vertical().show(ui, |ui| {
                    ui.heading("360° -> COLMAP cubemap rig");
                    ui.add_space(2.0);
                    ui.label(egui::RichText::new(format!("CLI: {}", self.cli_binary_str)).small().color(egui::Color32::GRAY));

                    ui.add_space(8.0);
                    ui.label("Project directory");
                    ui.horizontal(|ui| {
                        let mut s = self.project_dir.display().to_string();
                        let r = ui.text_edit_singleline(&mut s);
                        if r.changed() {
                            self.project_dir = PathBuf::from(s);
                        }
                        if ui.button("Browse…").clicked() {
                            if let Some(p) = rfd::FileDialog::new().pick_folder() {
                                self.project_dir = p;
                            }
                        }
                    });

                    ui.add_space(8.0);
                    ui.label("Source videos");
                    let vids = self.videos.clone();
                    for (i, v) in vids.iter().enumerate() {
                        ui.horizontal(|ui| {
                            ui.label(format!("{}. {}", i + 1, v.display()));
                            if ui.small_button("✕").clicked() {
                                self.videos.remove(i);
                            }
                        });
                    }
                    if ui.button("+ Add videos…").clicked() {
                        if let Some(files) = rfd::FileDialog::new()
                            .add_filter("Videos", &["mp4", "mov", "mkv", "avi", "webm", "m4v"])
                            .pick_files()
                        {
                            self.videos.extend(files);
                        }
                    }

                    ui.add_space(8.0);
                    ui.horizontal(|ui| {
                        ui.label("Start stage:");
                        egui::ComboBox::from_id_salt("start_stage")
                            .selected_text(self.start_stage.clone())
                            .show_ui(ui, |ui| {
                                for s in ["extract", "masks", "cubemap", "rig"] {
                                    ui.selectable_value(&mut self.start_stage, s.into(), s);
                                }
                            });
                    });
                    ui.checkbox(&mut self.skip_masks, "Skip masks stage");
                });
            });

        // Centre: stage option sections.
        egui::CentralPanel::default().show(ui, |ui| {
            egui::ScrollArea::vertical().show(ui, |ui| {
                ui.spacing_mut().item_spacing.y = 6.0;

                egui::CollapsingHeader::new("Stage 1 — Extract (sharpest/best frames)")
                    .default_open(true)
                    .show(ui, |ui| {
                        let e = &mut self.extract;
                        egui::Grid::new("extract_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
                            ui.label("Chunk size");        ui.add(egui::DragValue::new(&mut e.chunk_size).range(1..=10000)); ui.end_row();
                            ui.label("Scale width");       ui.add(egui::DragValue::new(&mut e.scale_width).range(0..=8000)); ui.end_row();
                            ui.label("Output format");    combo(ui, "extract_fmt", &["png", "jpg"], &mut e.output_format); ui.end_row();
                            ui.label("JPEG quality");      ui.add(egui::DragValue::new(&mut e.jpeg_quality).range(1..=100)); ui.end_row();
                            ui.label("Batch (frames/GPU)");ui.add(egui::DragValue::new(&mut e.batch).range(1..=64)); ui.end_row();
                            ui.label("Selection");         combo(ui, "extract_sel", &["chunk", "spread"], &mut e.selection); ui.end_row();
                            ui.label("Target count");     ui.add(egui::DragValue::new(&mut e.target_count).range(0..=100000)); ui.end_row();
                            ui.label("Min gap (frames)");  ui.add(egui::DragValue::new(&mut e.min_gap).range(0..=10000)); ui.end_row();
                            ui.label("Bad-frame action");  combo(ui, "extract_bfa", &["keep", "skip", "warn"], &mut e.bad_frame_action); ui.end_row();
                            ui.label("Motion mode");      combo(ui, "extract_mm", &["frame_diff", "off"], &mut e.motion_mode); ui.end_row();
                        });
                        ui.label(egui::RichText::new("Quality thresholds (0/255 bounds = no filter):").small());
                        let e = &mut self.extract;
                        egui::Grid::new("extract_q").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
                            ui.label("Min sharpness"); ui.add(egui::DragValue::new(&mut e.min_sharpness).speed(0.05).range(0.0..=1.0e9)); ui.end_row();
                            ui.label("Min exposure");  ui.add(egui::DragValue::new(&mut e.min_exposure).speed(1.0).range(0.0..=255.0)); ui.end_row();
                            ui.label("Max exposure");  ui.add(egui::DragValue::new(&mut e.max_exposure).speed(1.0).range(0.0..=255.0)); ui.end_row();
                            ui.label("Min contrast");  ui.add(egui::DragValue::new(&mut e.min_contrast).speed(1.0).range(0.0..=255.0)); ui.end_row();
                            ui.label("Min motion");    ui.add(egui::DragValue::new(&mut e.min_motion).speed(1.0).range(0.0..=255.0)); ui.end_row();
                            ui.label("Max motion");    ui.add(egui::DragValue::new(&mut e.max_motion).speed(1.0).range(0.0..=1.0e9)); ui.end_row();
                        });
                    });

                egui::CollapsingHeader::new("Stage 2 — Masks (YOLO segmentation; needs --features masks)")
                    .show(ui, |ui| {
                        let m = &mut self.masks;
                        egui::Grid::new("masks_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
                            ui.label("Classes");      ui.text_edit_singleline(&mut m.classes); ui.end_row();
                            ui.label("Model (ONNX)");  ui.horizontal(|ui| {
                                let mut s = m.model.display().to_string();
                                let r = ui.text_edit_singleline(&mut s);
                                if r.changed() { m.model = PathBuf::from(s); }
                                if ui.button("…").clicked() {
                                    if let Some(p) = rfd::FileDialog::new().add_filter("ONNX", &["onnx"]).pick_file() { m.model = p; }
                                }
                            });
                            ui.end_row();
                            ui.label("Device");       combo(ui, "mask_dev", &["cpu", "cuda", "rocm"], &mut m.device); ui.end_row();
                            ui.label("Conf");         ui.add(egui::DragValue::new(&mut m.conf).speed(0.01).range(0.0..=1.0).fixed_decimals(2)); ui.end_row();
                            ui.label("Imgsz");        ui.add(egui::DragValue::new(&mut m.imgsz).range(64..=4096)); ui.end_row();
                            ui.label("Dilate (px)");   ui.add(egui::DragValue::new(&mut m.dilate).range(0..=100)); ui.end_row();
                            ui.label("Projection");   combo(ui, "mask_proj", &["cubemap", "equirect"], &mut m.projection); ui.end_row();
                            ui.label("Face size");    ui.add(egui::DragValue::new(&mut m.face_size).range(64..=4096)); ui.end_row();
                            ui.label("Skip faces");   ui.text_edit_singleline(&mut m.skip_faces); ui.end_row();
                            ui.label("Tile batch");   ui.add(egui::DragValue::new(&mut m.tile_batch).range(1..=64)); ui.end_row();
                            ui.label("NMS IoU");      ui.add(egui::DragValue::new(&mut m.iou).speed(0.01).range(0.0..=1.0).fixed_decimals(2)); ui.end_row();
                        });
                        ui.checkbox(&mut m.empty_if_none, "Write empty (all-keep) mask when no objects detected");
                    });

                egui::CollapsingHeader::new("Stage 3 — Cubemap (equirect -> 6 faces)")
                    .show(ui, |ui| {
                        let c = &mut self.cubemap;
                        egui::Grid::new("cubemap_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
                            ui.label("Face size");     ui.add(egui::DragValue::new(&mut c.face_size).range(64..=4096)); ui.end_row();
                            ui.label("Skip faces");    ui.text_edit_singleline(&mut c.skip_faces); ui.end_row();
                            ui.label("JPEG quality");  ui.add(egui::DragValue::new(&mut c.jpeg_quality).range(1..=100)); ui.end_row();
                            ui.label("Batch (frames/GPU)"); ui.add(egui::DragValue::new(&mut c.batch).range(1..=32)); ui.end_row();
                        });
                    });

                egui::CollapsingHeader::new("Stage 4 — Rig config")
                    .show(ui, |ui| {
                        let r = &mut self.rig;
                        egui::Grid::new("rig_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
                            ui.label("Reference face"); combo(ui, "rig_ref", &["F", "R", "B", "L", "U", "D"], &mut r.ref_face); ui.end_row();
                            ui.label("Intrinsics");     combo(ui, "rig_int", &["linspace", "half"], &mut r.intrinsics_convention); ui.end_row();
                            ui.label("Link mode");      combo(ui, "rig_link", &["symlink", "hardlink", "copy"], &mut r.link_mode); ui.end_row();
                        });
                    });
            });
        });

        // Keep repainting while a run is in progress so the log streams live.
        if self.is_running() {
            ui.ctx().request_repaint();
        }
    }
}

pub fn launch() -> eframe::Result<()> {
    let options = eframe::NativeOptions {
        viewport: egui::ViewportBuilder::default()
            .with_inner_size([1100.0, 760.0])
            .with_min_inner_size([700.0, 500.0]),
        ..Default::default()
    };
    eframe::run_native(
        "pipeline_360 — 360° -> COLMAP cubemap rig",
        options,
        Box::new(|_cc| Ok(Box::<App>::default())),
    )
}