use eframe::egui;
use egui_plot::{Line, Plot, PlotPoints};
use std::collections::{HashMap, VecDeque};
use std::net::{TcpStream, UdpSocket};
use std::io::{Write, BufReader, BufRead};
use std::sync::{Arc, Mutex};
use std::thread;
use serde::{Deserialize, Serialize};
use chrono::Local;
use crossbeam_channel::{unbounded, Receiver, Sender};
use regex::Regex;

// --- Models ---

#[derive(Serialize, Deserialize, Clone, Debug)]
struct Signal {
    name: String,
    id: u32,
    #[serde(rename = "type")]
    sig_type: String,
}

#[derive(Deserialize)]
struct DiscoverResponse {
    #[serde(rename = "Signals")]
    signals: Vec<Signal>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
struct TreeItem {
    #[serde(rename = "Name")]
    name: String,
    #[serde(rename = "Class")]
    class: String,
    #[serde(rename = "Children")]
    children: Option<Vec<TreeItem>>,
    #[serde(rename = "Type")]
    sig_type: Option<String>,
    #[serde(rename = "Dimensions")]
    dimensions: Option<u8>,
    #[serde(rename = "Elements")]
    elements: Option<u32>,
}

#[derive(Clone)]
struct LogEntry {
    time: String,
    level: String,
    message: String,
}

struct TraceData {
    values: VecDeque<[f64; 2]>, 
}

struct SignalMetadata {
    names: Vec<String>,
    sig_type: String,
}

enum InternalEvent {
    Log(LogEntry),
    Discovery(Vec<Signal>),
    Tree(TreeItem),
    CommandResponse(String),
    NodeInfo(String),
    Connected,
    InternalLog(String),
    TraceRequested(String),
    ClearTrace(String),
    UdpStats(u64), 
}

// --- App State ---

struct ForcingDialog {
    signal_path: String,
    sig_type: String,
    dims: u8,
    elems: u32,
    value: String,
}

struct LogFilters {
    content_regex: String,
    show_debug: bool,
    show_info: bool,
    show_warning: bool,
    show_error: bool,
    paused: bool,
}

struct MarteDebugApp {
    #[allow(dead_code)]
    connected: bool,
    tcp_addr: String,
    log_addr: String,
    
    signals: Vec<Signal>,
    app_tree: Option<TreeItem>,
    
    traced_signals: Arc<Mutex<HashMap<String, TraceData>>>,
    id_to_meta: Arc<Mutex<HashMap<u32, SignalMetadata>>>,
    
    forced_signals: HashMap<String, String>,
    is_paused: bool,
    
    logs: VecDeque<LogEntry>,
    log_filters: LogFilters,
    
    // UI Panels
    show_left_panel: bool,
    show_right_panel: bool,
    show_bottom_panel: bool,
    
    selected_node: String,
    node_info: String,
    
    udp_packets: u64,
    
    forcing_dialog: Option<ForcingDialog>,
    
    tx_cmd: Sender<String>,
    rx_events: Receiver<InternalEvent>,
    internal_tx: Sender<InternalEvent>,
}

impl MarteDebugApp {
    fn new(_cc: &eframe::CreationContext<'_>) -> Self {
        let (tx_cmd, rx_cmd_internal) = unbounded::<String>();
        let (tx_events, rx_events) = unbounded::<InternalEvent>();
        let internal_tx = tx_events.clone();
        
        let tcp_addr = "127.0.0.1:8080".to_string();
        let log_addr = "127.0.0.1:8082".to_string();
        
        let id_to_meta = Arc::new(Mutex::new(HashMap::new()));
        let traced_signals = Arc::new(Mutex::new(HashMap::new()));
        
        let id_to_meta_clone = id_to_meta.clone();
        let traced_signals_clone = traced_signals.clone();

        let tx_events_c = tx_events.clone();
        thread::spawn(move || {
            tcp_command_worker(tcp_addr, rx_cmd_internal, tx_events_c);
        });

        let tx_events_log = tx_events.clone();
        thread::spawn(move || {
            tcp_log_worker(log_addr, tx_events_log);
        });

        let tx_events_udp = tx_events.clone();
        thread::spawn(move || {
            udp_worker(8081, id_to_meta_clone, traced_signals_clone, tx_events_udp);
        });

        Self {
            connected: false,
            tcp_addr: "127.0.0.1:8080".to_string(),
            log_addr: "127.0.0.1:8082".to_string(),
            signals: Vec::new(),
            app_tree: None,
            id_to_meta,
            traced_signals,
            forced_signals: HashMap::new(),
            is_paused: false,
            logs: VecDeque::with_capacity(2000),
            log_filters: LogFilters {
                content_regex: "".to_string(),
                show_debug: true,
                show_info: true,
                show_warning: true,
                show_error: true,
                paused: false,
            },
            show_left_panel: true,
            show_right_panel: true,
            show_bottom_panel: true,
            selected_node: "".to_string(),
            node_info: "".to_string(),
            udp_packets: 0,
            forcing_dialog: None,
            tx_cmd,
            rx_events,
            internal_tx,
        }
    }

    fn render_tree(&mut self, ui: &mut egui::Ui, item: &TreeItem, path: String) {
        let current_path = if path.is_empty() { 
            item.name.clone() 
        } else if path == "Root" {
            item.name.clone()
        } else { 
            format!("{}.{}", path, item.name) 
        };
        
        if let Some(children) = &item.children {
            let header = egui::CollapsingHeader::new(format!("{} [{}]", item.name, item.class))
                .id_salt(&current_path);
            
            header.show(ui, |ui| {
                ui.horizontal(|ui| {
                    if ui.selectable_label(self.selected_node == current_path, "ℹ Info").clicked() {
                        self.selected_node = current_path.clone();
                        let _ = self.tx_cmd.send(format!("INFO {}", current_path));
                    }
                });
                for child in children {
                    self.render_tree(ui, child, current_path.clone());
                }
            });
        } else {
            ui.horizontal(|ui| {
                if ui.selectable_label(self.selected_node == current_path, format!("{} [{}]", item.name, item.class)).clicked() {
                    self.selected_node = current_path.clone();
                    let _ = self.tx_cmd.send(format!("INFO {}", current_path));
                }
                if item.class.contains("Signal") {
                    if ui.button("📈 Trace").clicked() {
                        let _ = self.tx_cmd.send(format!("TRACE {} 1", current_path));
                        let _ = self.internal_tx.send(InternalEvent::TraceRequested(current_path.clone()));
                    }
                    if ui.button("⚡ Force").clicked() {
                        self.forcing_dialog = Some(ForcingDialog {
                            signal_path: current_path.clone(),
                            sig_type: item.sig_type.clone().unwrap_or_else(|| "Unknown".to_string()),
                            dims: item.dimensions.unwrap_or(0),
                            elems: item.elements.unwrap_or(1),
                            value: "".to_string(),
                        });
                    }
                }
            });
        }
    }
}

fn tcp_command_worker(addr: String, rx_cmd: Receiver<String>, tx_events: Sender<InternalEvent>) {
    loop {
        if let Ok(mut stream) = TcpStream::connect(&addr) {
            let _ = stream.set_nodelay(true);
            let mut reader = BufReader::new(stream.try_clone().unwrap());
            let _ = tx_events.send(InternalEvent::Connected);
            
            let tx_events_inner = tx_events.clone();
            thread::spawn(move || {
                let mut line = String::new();
                let mut json_acc = String::new();
                let mut in_json = false;

                while reader.read_line(&mut line).is_ok() {
                    let trimmed = line.trim();
                    if trimmed.is_empty() { line.clear(); continue; }
                    
                    if !in_json && trimmed.starts_with("{") {
                        in_json = true;
                        json_acc.clear();
                    }

                    if in_json {
                        json_acc.push_str(trimmed);
                        if trimmed == "OK DISCOVER" {
                            in_json = false;
                            let json_clean = json_acc.trim_end_matches("OK DISCOVER").trim();
                            match serde_json::from_str::<DiscoverResponse>(json_clean) {
                                Ok(resp) => { let _ = tx_events_inner.send(InternalEvent::Discovery(resp.signals)); }
                                Err(e) => { let _ = tx_events_inner.send(InternalEvent::InternalLog(format!("JSON Parse Error (Discover): {}", e))); }
                            }
                            json_acc.clear();
                        } else if trimmed == "OK TREE" {
                            in_json = false;
                            let json_clean = json_acc.trim_end_matches("OK TREE").trim();
                            match serde_json::from_str::<TreeItem>(json_clean) {
                                Ok(resp) => { let _ = tx_events_inner.send(InternalEvent::Tree(resp)); }
                                Err(e) => { let _ = tx_events_inner.send(InternalEvent::InternalLog(format!("JSON Parse Error (Tree): {}", e))); }
                            }
                            json_acc.clear();
                        } else if trimmed == "OK INFO" {
                            in_json = false;
                            let json_clean = json_acc.trim_end_matches("OK INFO").trim();
                            let _ = tx_events_inner.send(InternalEvent::NodeInfo(json_clean.to_string()));
                            json_acc.clear();
                        }
                    } else {
                        let _ = tx_events_inner.send(InternalEvent::CommandResponse(trimmed.to_string()));
                    }
                    line.clear();
                }
            });

            while let Ok(cmd) = rx_cmd.recv() {
                if stream.write_all(format!("{}\n", cmd).as_bytes()).is_err() {
                    break;
                }
            }
        }
        thread::sleep(std::time::Duration::from_secs(2));
    }
}

fn tcp_log_worker(addr: String, tx_events: Sender<InternalEvent>) {
    loop {
        if let Ok(stream) = TcpStream::connect(&addr) {
            let mut reader = BufReader::new(stream);
            let mut line = String::new();
            while reader.read_line(&mut line).is_ok() {
                let trimmed = line.trim();
                if trimmed.starts_with("LOG ") {
                    let parts: Vec<&str> = trimmed[4..].splitn(2, ' ').collect();
                    if parts.len() == 2 {
                        let _ = tx_events.send(InternalEvent::Log(LogEntry {
                            time: Local::now().format("%H:%M:%S%.3f").to_string(),
                            level: parts[0].to_string(),
                            message: parts[1].to_string(),
                        }));
                    }
                }
                line.clear();
            }
        }
        thread::sleep(std::time::Duration::from_secs(2));
    }
}

fn udp_worker(port: u16, id_to_meta: Arc<Mutex<HashMap<u32, SignalMetadata>>>, traced_data: Arc<Mutex<HashMap<String, TraceData>>>, tx_events: Sender<InternalEvent>) {
    if let Ok(socket) = UdpSocket::bind(format!("0.0.0.0:{}", port)) {
        let mut buf = [0u8; 4096];
        let start_time = std::time::Instant::now();
        let mut total_packets = 0u64;

        loop {
            if let Ok(n) = socket.recv(&mut buf) {
                total_packets += 1;
                if (total_packets % 100) == 0 {
                    let _ = tx_events.send(InternalEvent::UdpStats(total_packets));
                }

                if n < 20 { continue; }
                
                let mut magic_buf = [0u8; 4]; magic_buf.copy_from_slice(&buf[0..4]);
                if u32::from_le_bytes(magic_buf) != 0xDA7A57AD { continue; }

                let mut count_buf = [0u8; 4]; count_buf.copy_from_slice(&buf[16..20]);
                let count = u32::from_le_bytes(count_buf);
                
                let mut offset = 20;
                let now = start_time.elapsed().as_secs_f64();

                let metas = id_to_meta.lock().unwrap();
                let mut data_map = traced_data.lock().unwrap();

                for _ in 0..count {
                    if offset + 8 > n { break; }
                    let mut id_buf = [0u8; 4]; id_buf.copy_from_slice(&buf[offset..offset+4]);
                    let id = u32::from_le_bytes(id_buf);
                    let mut size_buf = [0u8; 4]; size_buf.copy_from_slice(&buf[offset+4..offset+8]);
                    let size = u32::from_le_bytes(size_buf);
                    offset += 8;
                    if offset + size as usize > n { break; }
                    
                    let data_slice = &buf[offset..offset + size as usize];

                    if let Some(meta) = metas.get(&id) {
                        let t = meta.sig_type.as_str();
                        let val = match size {
                            1 => { if t.contains('u') { data_slice[0] as f64 } else { (data_slice[0] as i8) as f64 } },
                            2 => {
                                let mut b = [0u8; 2]; b.copy_from_slice(data_slice);
                                if t.contains('u') { u16::from_le_bytes(b) as f64 } else { i16::from_le_bytes(b) as f64 }
                            },
                            4 => {
                                let mut b = [0u8; 4]; b.copy_from_slice(data_slice);
                                if t.contains("float") { f32::from_le_bytes(b) as f64 }
                                else if t.contains('u') { u32::from_le_bytes(b) as f64 }
                                else { i32::from_le_bytes(b) as f64 }
                            },
                            8 => {
                                let mut b = [0u8; 8]; b.copy_from_slice(data_slice);
                                if t.contains("float") { f64::from_le_bytes(b) }
                                else if t.contains('u') { u64::from_le_bytes(b) as f64 }
                                else { i64::from_le_bytes(b) as f64 }
                            },
                            _ => 0.0,
                        };

                        for name in &meta.names {
                            if let Some(entry) = data_map.get_mut(name) {
                                entry.values.push_back([now, val]);
                                if entry.values.len() > 2000 { entry.values.pop_front(); }
                            }
                        }
                    }
                    offset += size as usize;
                }
            }
        }
    }
}

impl eframe::App for MarteDebugApp {
    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
        while let Ok(event) = self.rx_events.try_recv() {
            match event {
                InternalEvent::Log(log) => {
                    if !self.log_filters.paused {
                        self.logs.push_back(log);
                        if self.logs.len() > 2000 { self.logs.pop_front(); }
                    }
                }
                InternalEvent::InternalLog(msg) => {
                    self.logs.push_back(LogEntry {
                        time: Local::now().format("%H:%M:%S").to_string(),
                        level: "GUI_ERROR".to_string(),
                        message: msg,
                    });
                }
                InternalEvent::Discovery(signals) => {
                    let mut metas = self.id_to_meta.lock().unwrap();
                    metas.clear();
                    for s in &signals {
                        let meta = metas.entry(s.id).or_insert_with(|| SignalMetadata { names: Vec::new(), sig_type: s.sig_type.clone() });
                        if !meta.names.contains(&s.name) { meta.names.push(s.name.clone()); }
                    }
                    self.signals = signals;
                }
                InternalEvent::Tree(tree) => {
                    self.app_tree = Some(tree);
                }
                InternalEvent::NodeInfo(info) => {
                    self.node_info = info;
                }
                InternalEvent::TraceRequested(name) => {
                    let mut data_map = self.traced_signals.lock().unwrap();
                    data_map.entry(name).or_insert_with(|| TraceData { values: VecDeque::with_capacity(2000) });
                }
                InternalEvent::ClearTrace(name) => {
                    let mut data_map = self.traced_signals.lock().unwrap();
                    data_map.remove(&name);
                }
                InternalEvent::CommandResponse(resp) => {
                    self.logs.push_back(LogEntry {
                        time: Local::now().format("%H:%M:%S").to_string(),
                        level: "CMD_RESP".to_string(),
                        message: resp,
                    });
                }
                InternalEvent::UdpStats(count) => {
                    self.udp_packets = count;
                }
                InternalEvent::Connected => {
                    let _ = self.tx_cmd.send("TREE".to_string());
                    let _ = self.tx_cmd.send("DISCOVER".to_string());
                }
            }
        }

        if let Some(dialog) = &mut self.forcing_dialog {
            let mut close = false;
            egui::Window::new("Force Signal").collapsible(false).resizable(false).show(ctx, |ui| {
                ui.label(format!("Signal: {}", dialog.signal_path));
                ui.label(format!("Type: {} (Dims: {}, Elems: {})", dialog.sig_type, dialog.dims, dialog.elems));
                ui.horizontal(|ui| {
                    ui.label("Value:");
                    ui.text_edit_singleline(&mut dialog.value);
                });
                ui.horizontal(|ui| {
                    if ui.button("Apply Force").clicked() {
                        let _ = self.tx_cmd.send(format!("FORCE {} {}", dialog.signal_path, dialog.value));
                        self.forced_signals.insert(dialog.signal_path.clone(), dialog.value.clone());
                        close = true;
                    }
                    if ui.button("Cancel").clicked() { close = true; }
                });
            });
            if close { self.forcing_dialog = None; }
        }

        egui::TopBottomPanel::top("top_panel").show(ctx, |ui| {
            ui.horizontal(|ui| {
                ui.toggle_value(&mut self.show_left_panel, "🗂 Tree");
                ui.toggle_value(&mut self.show_right_panel, "⚙ Debug");
                ui.toggle_value(&mut self.show_bottom_panel, "📜 Logs");
                ui.separator();
                ui.heading("MARTe2 Debug Explorer");
                ui.separator();
                if ui.button("🔄 Refresh").clicked() {
                    let _ = self.tx_cmd.send("TREE".to_string());
                    let _ = self.tx_cmd.send("DISCOVER".to_string());
                }
                ui.separator();
                if self.is_paused {
                    if ui.button("▶ Resume").clicked() {
                        let _ = self.tx_cmd.send("RESUME".to_string());
                        self.is_paused = false;
                    }
                } else {
                    if ui.button("⏸ Pause").clicked() {
                        let _ = self.tx_cmd.send("PAUSE".to_string());
                        self.is_paused = true;
                    }
                }
                ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                    ui.label(format!("UDP Packets: {}", self.udp_packets));
                });
            });
        });

        if self.show_bottom_panel {
            egui::TopBottomPanel::bottom("log_panel").resizable(true).default_height(200.0).show(ctx, |ui| {
                ui.horizontal(|ui| {
                    ui.heading("System Logs");
                    ui.separator();
                    ui.label("Filter:");
                    ui.text_edit_singleline(&mut self.log_filters.content_regex);
                    ui.checkbox(&mut self.log_filters.show_debug, "Debug");
                    ui.checkbox(&mut self.log_filters.show_info, "Info");
                    ui.checkbox(&mut self.log_filters.show_warning, "Warn");
                    ui.checkbox(&mut self.log_filters.show_error, "Error");
                    ui.separator();
                    ui.toggle_value(&mut self.log_filters.paused, "⏸ Pause Logs");
                    ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                        if ui.button("🗑 Clear").clicked() { self.logs.clear(); }
                    });
                });
                ui.separator();
                
                let regex = if !self.log_filters.content_regex.is_empty() { Regex::new(&self.log_filters.content_regex).ok() } else { None };

                egui::ScrollArea::vertical()
                    .stick_to_bottom(true)
                    .auto_shrink([false, false])
                    .show(ui, |ui| {
                        for log in &self.logs {
                            let show = match log.level.as_str() {
                                "Debug" => self.log_filters.show_debug,
                                "Information" => self.log_filters.show_info,
                                "Warning" => self.log_filters.show_warning,
                                "FatalError" | "OSError" | "ParametersError" => self.log_filters.show_error,
                                _ => true,
                            };
                            if !show { continue; }
                            if let Some(re) = &regex { if !re.is_match(&log.message) && !re.is_match(&log.level) { continue; } }

                            let color = match log.level.as_str() {
                                "FatalError" | "OSError" | "ParametersError" => egui::Color32::from_rgb(255, 100, 100),
                                "Warning" => egui::Color32::from_rgb(255, 255, 100),
                                "Information" => egui::Color32::from_rgb(100, 255, 100),
                                "Debug" => egui::Color32::from_rgb(100, 100, 255),
                                "GUI_ERROR" => egui::Color32::from_rgb(255, 50, 255),
                                "CMD_RESP" => egui::Color32::from_rgb(255, 255, 255),
                                _ => egui::Color32::WHITE,
                            };
                            
                            ui.scope(|ui| {
                                ui.horizontal(|ui| {
                                    ui.label(egui::RichText::new(&log.time).monospace().color(egui::Color32::GRAY));
                                    ui.label(egui::RichText::new(format!("[{}]", log.level)).color(color).strong());
                                    ui.add(egui::Label::new(&log.message).wrap());
                                    ui.allocate_space(egui::vec2(ui.available_width(), 0.0));
                                });
                            });
                        }
                    });
            });
        }

        if self.show_left_panel {
            egui::SidePanel::left("signals_panel").resizable(true).width_range(300.0..=600.0).show(ctx, |ui| {
                ui.heading("Application Tree");
                ui.separator();
                egui::ScrollArea::vertical().id_salt("tree_scroll").show(ui, |ui| {
                    if let Some(tree) = self.app_tree.clone() {
                        self.render_tree(ui, &tree, "".to_string());
                    } else {
                        ui.label("Connecting to server...");
                    }
                });
                
                ui.separator();
                ui.heading("Node Information");
                egui::ScrollArea::vertical().id_salt("info_scroll").show(ui, |ui| {
                    if self.node_info.is_empty() {
                        ui.label("Click 'Info' on a node to view details");
                    } else {
                        ui.add(egui::Label::new(egui::RichText::new(&self.node_info).monospace()).wrap());
                    }
                });
            });
        }

        if self.show_right_panel {
            egui::SidePanel::right("debug_panel").resizable(true).width_range(200.0..=400.0).show(ctx, |ui| {
                ui.heading("Active Controls");
                ui.separator();
                ui.label(egui::RichText::new("Forced Signals").strong());
                egui::ScrollArea::vertical().id_salt("forced_scroll").show(ui, |ui| {
                    let mut to_update = None;
                    let mut to_remove = None;
                    for (path, val) in &self.forced_signals {
                        ui.horizontal(|ui| {
                            if ui.selectable_label(false, format!("{}: {}", path, val)).clicked() {
                                to_update = Some(path.clone());
                            }
                            if ui.button("❌").clicked() {
                                let _ = self.tx_cmd.send(format!("UNFORCE {}", path));
                                to_remove = Some(path.clone());
                            }
                        });
                    }
                    if let Some(p) = to_remove { self.forced_signals.remove(&p); }
                    if let Some(p) = to_update {
                        self.forcing_dialog = Some(ForcingDialog {
                            signal_path: p.clone(), sig_type: "Unknown".to_string(), dims: 0, elems: 1,
                            value: self.forced_signals.get(&p).unwrap().clone(),
                        });
                    }
                });

                ui.separator();
                ui.label(egui::RichText::new("Traced Signals").strong());
                egui::ScrollArea::vertical().id_salt("traced_scroll").show(ui, |ui| {
                    let mut names: Vec<_> = {
                        let data_map = self.traced_signals.lock().unwrap();
                        data_map.keys().cloned().collect()
                    };
                    names.sort();
                    for key in names {
                        ui.horizontal(|ui| {
                            ui.label(&key);
                            if ui.button("❌").clicked() {
                                let _ = self.tx_cmd.send(format!("TRACE {} 0", key));
                                let _ = self.internal_tx.send(InternalEvent::ClearTrace(key));
                            }
                        });
                    }
                });
            });
        }

        egui::CentralPanel::default().show(ctx, |ui| {
            ui.heading("Oscilloscope");
            let plot = Plot::new("traces_plot")
                .legend(egui_plot::Legend::default())
                .auto_bounds_x()
                .auto_bounds_y()
                .y_axis_min_width(4.0);
            
            plot.show(ui, |plot_ui| {
                let data_map = self.traced_signals.lock().unwrap();
                for (name, data) in data_map.iter() {
                    let points: PlotPoints = data.values.iter().cloned().collect();
                    plot_ui.line(Line::new(points).name(name));
                }
            });
        });

        ctx.request_repaint_after(std::time::Duration::from_millis(16));
    }
}

fn main() -> Result<(), eframe::Error> {
    let options = eframe::NativeOptions {
        viewport: egui::ViewportBuilder::default().with_inner_size([1280.0, 800.0]),
        ..Default::default()
    };
    eframe::run_native(
        "MARTe2 Debug Explorer",
        options,
        Box::new(|cc| Ok(Box::new(MarteDebugApp::new(cc)))),
    )
}