// Copyright 2025 David Weikersdorfer

use eyre::{bail, ContextCompat, Result};
use nodo::{
    codelet::{ScheduleId, Transition},
    prelude::{ParameterDataType, ParameterValue, ParameterWithPropertiesSet},
};
use nodo_runtime::proto::nodo as nodo_pb;
use std::{cell::RefCell, cmp::Ordering, collections::HashMap, iter, rc::Rc};

/// Stores a hierarchical data structure worker -> sequence -> node -> parameter
#[derive(Default)]
pub struct NodoAppTree {
    pub last_report_app_meta: Option<nodo_pb::AppMeta>,
    pub schedules: HashMap<String, NodoAppTreeSchedule>,
    pub monitors: Vec<(String, nodo_pb::MonitorMeta, nodo_pb::Monitor)>,

    pub step_duration: f32,
}

impl NodoAppTree {
    pub fn update_step_duration(&mut self) {
        self.step_duration = 0.0;

        for (_, sched) in self.schedules.iter_mut() {
            sched.step_duration = 0.0;

            for (_, seq) in sched.seqs.iter_mut() {
                seq.step_duration = 0.0;

                for (_, node) in seq.nodes.iter() {
                    let step_transition = &node.statistics.transitions[Transition::Step.index()];
                    let delta = step_transition.duration.context("TODO").unwrap().total;
                    seq.step_duration += delta;
                }

                sched.step_duration += seq.step_duration;
            }

            self.step_duration += sched.step_duration;
        }
    }
}

pub struct NodoAppTreeSchedule {
    pub id: ScheduleId,
    pub name: String,

    pub target_period: Option<f32>,
    pub status: nodo_pb::ScheduleStatus,

    pub is_expanded: bool,
    pub is_selected: bool,
    pub tree_row_idx: RefCell<Option<usize>>,

    pub seqs: HashMap<String, NodoAppTreeSequence>,

    pub step_duration: f32,
}

pub struct NodoAppTreeSequence {
    pub name: String,

    pub is_expanded: bool,
    pub is_selected: bool,
    pub tree_row_idx: RefCell<Option<usize>>,

    pub nodes: HashMap<String, NodoAppTreeNode>,

    pub step_duration: f32,
}

pub struct NodoAppTreeNode {
    pub info: nodo_pb::NodeMeta,

    pub status: nodo_pb::NodeStatus,
    pub statistics: nodo_pb::NodeStatistics,
    pub signals: Vec<(String, nodo_pb::NodeSignal)>,

    pub parameters: HashMap<String, NodoAppTreeParameter>,

    pub is_selected: bool,
    pub tree_row_idx: RefCell<Option<usize>>,
}

impl NodoAppTreeNode {
    pub fn step_duration(&self) -> f32 {
        self.statistics.transitions[Transition::Step.index()]
            .duration
            .unwrap()
            .total
    }
}

pub struct NodoAppTreeParameter {
    pub name: String,
    pub dtype: ParameterDataType,
    pub is_mutable: bool,
    pub value: ParameterValue,
}

impl NodoAppTree {
    pub fn update(&mut self, report: nodo_pb::Report) -> Result<()> {
        if let Some(app_meta) = report.meta {
            self.last_report_app_meta = Some(app_meta);
        }
        if self.last_report_app_meta.is_none() {
            bail!("app info not available");
        }

        let (
            Some(nodo_pb::AppMeta {
                nodes: node_meta,
                monitors: monitor_meta,
                schedules: schedule_meta,
            }),
            Some(nodo_pb::AppStatus {
                nodes: status,
                schedules: schedule_status,
            }),
            Some(nodo_pb::AppStatistics { nodes: statistics }),
            Some(nodo_pb::AppSignals { nodes: signals }),
        ) = (
            &self.last_report_app_meta,
            report.status,
            report.statistics,
            report.signals,
        )
        else {
            bail!("invalid report");
        };

        if node_meta.len() != status.len()
            || status.len() != statistics.len()
            || statistics.len() != signals.len()
        {
            bail!("invalid report");
        }

        if let Some(monitors) = report.monitors.map(|am| am.monitors) {
            if monitor_meta.len() != monitors.len() {
                bail!("invalid report");
            }

            self.monitors = monitor_meta
                .iter()
                .zip(monitors.into_iter())
                .map(|(meta, mon)| {
                    let node_name = node_meta
                        .iter()
                        .find(|i| i.id == meta.node_id)
                        .map(|i| i.node_name.clone())
                        .unwrap_or_else(|| "(unknown)".into());
                    (node_name, meta.clone(), mon)
                })
                .collect();
        }

        // delete nodes which don't match
        Self::retain_nodes(&mut self.schedules, |nid| {
            node_meta.iter().any(|n| n.id == nid)
        });

        // sync schedules
        self.schedules
            .retain(|needle, _| schedule_meta.iter().any(|sched| &sched.name == needle));
        for sched in schedule_meta.iter() {
            if !self.schedules.contains_key(&sched.name) {
                self.schedules.insert(
                    sched.name.clone(),
                    NodoAppTreeSchedule {
                        id: ScheduleId(sched.id as usize),
                        name: sched.name.clone(),
                        target_period: sched.period,
                        status: nodo_pb::ScheduleStatus::default(),
                        is_expanded: true,
                        is_selected: false,
                        tree_row_idx: RefCell::new(None),
                        seqs: HashMap::new(),
                        step_duration: 0.,
                    },
                );
            }
        }

        for status in schedule_status.iter() {
            match self
                .schedules
                .values_mut()
                .find(|v| v.id.0 == status.id as usize)
            {
                Some(sched) => {
                    sched.status = status.clone();
                }
                None => {
                    // mal-formed report
                }
            }
        }

        // update existing nodes / create new nodes
        for (info, (status, (statistics, nodo_pb::NodeSignals { signals }))) in
            node_meta.iter().zip(
                status
                    .into_iter()
                    .zip(statistics.into_iter().zip(signals.into_iter())),
            )
        {
            let Some(worker) = self.schedules.get_mut(&info.schedule_name) else {
                // FIXME this indicates an invalid report
                continue;
            };

            let sequence = worker
                .seqs
                .entry(info.sequence_name.clone())
                .or_insert_with_key(|key| NodoAppTreeSequence {
                    name: key.into(),
                    is_expanded: true,
                    is_selected: false,
                    tree_row_idx: RefCell::new(None),
                    nodes: HashMap::default(),
                    step_duration: 0.0,
                });

            let named_signals = info
                .signal_names
                .iter()
                .cloned()
                .zip(signals.iter().cloned())
                .collect::<Vec<_>>();

            let node = sequence
                .nodes
                .entry(info.node_name.clone())
                .or_insert_with_key(|_| NodoAppTreeNode {
                    info: info.clone(),
                    status: status.clone(),
                    statistics: statistics.clone(),
                    signals: named_signals.clone(),
                    parameters: HashMap::default(),
                    is_selected: false,
                    tree_row_idx: RefCell::new(None),
                });

            node.status = status;
            node.statistics = statistics;
            node.signals = named_signals;
        }

        self.update_step_duration();

        Ok(())
    }

    pub fn update_config_parameters(&mut self, items: ParameterWithPropertiesSet<String, String>) {
        for (_, worker) in self.schedules.iter_mut() {
            for (_, seq) in worker.seqs.iter_mut() {
                for (node_name, node) in seq.nodes.iter_mut() {
                    for (key, (props, value)) in items.iter() {
                        if node_name == key.node() {
                            let tp =
                                node.parameters
                                    .entry(key.param().clone())
                                    .or_insert_with(|| NodoAppTreeParameter {
                                        name: key.param().clone(),
                                        dtype: props.dtype,
                                        is_mutable: props.is_mutable,
                                        value: value.clone(),
                                    });

                            tp.dtype = value.dtype();
                            tp.value = value.clone();
                        }
                    }
                }
            }
        }
    }

    /// Iterate over all workers
    pub fn iter_workers(&self) -> impl Iterator<Item = &NodoAppTreeSchedule> {
        self.schedules.iter().map(|(_, v)| v)
    }

    /// Iterate over all workers
    pub fn iter_workers_mut(&mut self) -> impl Iterator<Item = &mut NodoAppTreeSchedule> {
        self.schedules.iter_mut().map(|(_, v)| v)
    }

    /// Iterate over all sequences
    pub fn iter_sequences(&self) -> impl Iterator<Item = (&str, &NodoAppTreeSequence)> {
        self.schedules
            .iter()
            .flat_map(|(sched_name, w)| w.seqs.iter().map(|(_, seq)| (sched_name.as_str(), seq)))
    }

    /// Iterate over all sequences
    pub fn iter_sequences_mut(&mut self) -> impl Iterator<Item = (&str, &mut NodoAppTreeSequence)> {
        self.schedules.iter_mut().flat_map(|(sched_name, w)| {
            w.seqs
                .iter_mut()
                .map(move |(_, seq)| (sched_name.as_str(), seq))
        })
    }

    /// Iterate over all nodes
    pub fn iter_nodes(
        &self,
    ) -> impl Iterator<Item = (&NodoAppTreeSchedule, &NodoAppTreeSequence, &NodoAppTreeNode)> {
        self.schedules.iter().flat_map(|(_, w)| {
            w.seqs
                .iter()
                .flat_map(move |(_, seq)| seq.nodes.iter().map(move |(_, node)| (w, seq, node)))
        })
    }

    /// Iterate over all config parameters
    pub fn iter_parameters(
        &self,
    ) -> impl Iterator<
        Item = (
            &NodoAppTreeSchedule,
            &NodoAppTreeSequence,
            &NodoAppTreeNode,
            &NodoAppTreeParameter,
        ),
    > {
        self.schedules.iter().flat_map(|(_, w)| {
            w.seqs.iter().flat_map(move |(_, seq)| {
                seq.nodes.iter().flat_map(move |(_, node)| {
                    node.parameters
                        .iter()
                        .map(move |(_, param)| (w, seq, node, param))
                })
            })
        })
    }

    /// Iterate over all signals
    pub fn iter_signals(
        &self,
    ) -> impl Iterator<
        Item = (
            &NodoAppTreeSchedule,
            &NodoAppTreeSequence,
            &NodoAppTreeNode,
            &(String, nodo_pb::NodeSignal),
        ),
    > {
        self.schedules.iter().flat_map(|(_, w)| {
            w.seqs.iter().flat_map(move |(_, seq)| {
                seq.nodes.iter().flat_map(move |(_, node)| {
                    node.signals.iter().map(move |entry| (w, seq, node, entry))
                })
            })
        })
    }

    /// Iterate over all tree items
    pub fn iter<'a>(&'a self) -> impl Iterator<Item = AppTreeItem<'a>> {
        self.schedules.iter().flat_map(|(_, w)| {
            iter::once(AppTreeItem::Worker(w)).chain(w.seqs.iter().flat_map(|(_, seq)| {
                iter::once(AppTreeItem::Sequence(seq))
                    .chain(seq.nodes.iter().map(|(_, node)| AppTreeItem::Node(node)))
            }))
        })
    }

    /// Iterate over all tree items in given order
    pub fn tree_iter<'a, F1, F2>(
        &'a self,
        cmp: F1,
        recurse: F2,
    ) -> impl Iterator<Item = (TreeLocation, AppTreeItem<'a>)>
    where
        F1: Fn(&AppTreeItem<'a>, &AppTreeItem<'a>) -> Ordering + Copy,
        F2: Fn(&AppTreeItem<'a>) -> bool + Copy,
    {
        let mut schedules: Vec<_> = self
            .schedules
            .iter()
            .map(|(_, v)| AppTreeItem::Worker(v))
            .collect();
        schedules.sort_by(|a, b| cmp(a, b));
        let sched_count = schedules.len();

        let tree_row_idx = Rc::new(RefCell::new(0));

        schedules
            .into_iter()
            .enumerate()
            .flat_map(move |(group_idx, sched)| {
                let AppTreeItem::Worker(sched) = sched else {
                    unreachable!()
                };

                *sched.tree_row_idx.borrow_mut() = Some(*tree_row_idx.borrow());
                *tree_row_idx.borrow_mut() += 1;

                let head = (
                    TreeLocation {
                        level: TreeLocationLevel::Node(0),
                        pos: TreeLocationPos::from_index_count(group_idx, sched_count),
                        is_expanded: recurse(&AppTreeItem::Worker(sched)),
                    },
                    AppTreeItem::Worker(sched),
                );

                let mut seqs = Vec::new();
                if head.0.is_expanded {
                    seqs.extend(sched.seqs.iter().map(|(_, v)| AppTreeItem::Sequence(v)));
                }
                seqs.sort_by(|a, b| cmp(a, b));
                let seq_count = seqs.len();

                iter::once(head).chain(seqs.into_iter().enumerate().flat_map({
                    let tree_row_idx = tree_row_idx.clone();
                    move |(group_idx, seq)| {
                        let AppTreeItem::Sequence(seq) = seq else {
                            unreachable!()
                        };

                        *seq.tree_row_idx.borrow_mut() = Some(*tree_row_idx.borrow());
                        *tree_row_idx.borrow_mut() += 1;

                        let head = (
                            TreeLocation {
                                level: TreeLocationLevel::Node(1),
                                pos: TreeLocationPos::from_index_count(group_idx, seq_count),
                                is_expanded: recurse(&AppTreeItem::Sequence(seq)),
                            },
                            AppTreeItem::Sequence(seq),
                        );

                        let mut nodes = Vec::new();
                        if head.0.is_expanded {
                            nodes.extend(seq.nodes.iter().map(|(_, v)| AppTreeItem::Node(v)));
                        }
                        nodes.sort_by(|a, b| cmp(a, b));
                        let node_count = nodes.len();

                        iter::once(head).chain(nodes.into_iter().enumerate().map({
                            let tree_row_idx = tree_row_idx.clone();
                            move |(group_idx, node)| {
                                let AppTreeItem::Node(node) = node else {
                                    unreachable!()
                                };

                                *node.tree_row_idx.borrow_mut() = Some(*tree_row_idx.borrow());
                                *tree_row_idx.borrow_mut() += 1;

                                (
                                    TreeLocation {
                                        level: TreeLocationLevel::Leaf,
                                        pos: TreeLocationPos::from_index_count(
                                            group_idx, node_count,
                                        ),
                                        is_expanded: true,
                                    },
                                    AppTreeItem::Node(node),
                                )
                            }
                        }))
                    }
                }))
            })
    }

    /// Selects the i-th tree entry based on the order used by tree_iter
    pub fn tree_entry_select(&mut self, selection: Option<usize>) {
        for (_, sched) in self.schedules.iter_mut() {
            sched.is_selected = *sched.tree_row_idx.borrow() == selection;

            for (_, seq) in sched.seqs.iter_mut() {
                seq.is_selected = *seq.tree_row_idx.borrow() == selection;

                for (_, node) in seq.nodes.iter_mut() {
                    node.is_selected = *node.tree_row_idx.borrow() == selection;
                }
            }
        }
    }

    /// Toggles the expand flag for selected nodes
    pub fn on_toggle_expand_selected(&mut self) {
        for (_, sched) in self.schedules.iter_mut() {
            if sched.is_selected {
                sched.is_expanded = !sched.is_expanded;
            }

            for (_, seq) in sched.seqs.iter_mut() {
                if seq.is_selected {
                    seq.is_expanded = !seq.is_expanded;
                }
            }
        }
    }

    /// Deletes all nodes which do not meet the criteria
    fn retain_nodes<F>(workers: &mut HashMap<String, NodoAppTreeSchedule>, f: F)
    where
        F: Fn(u32) -> bool,
    {
        for (_, wrk) in workers.iter_mut() {
            for (_, seq) in wrk.seqs.iter_mut() {
                seq.nodes.retain(|_, n| f(n.info.id))
            }
        }
    }
}

pub enum AppTreeItem<'a> {
    Worker(&'a NodoAppTreeSchedule),
    Sequence(&'a NodoAppTreeSequence),
    Node(&'a NodoAppTreeNode),
}

pub struct TreeLocation {
    pub level: TreeLocationLevel,
    pub pos: TreeLocationPos,
    pub is_expanded: bool,
}

pub enum TreeLocationLevel {
    Node(usize),
    Leaf,
}

pub enum TreeLocationPos {
    First,
    Mid,
    Last,
}

impl TreeLocationPos {
    pub fn from_index_count(i: usize, n: usize) -> Self {
        if i == 0 {
            TreeLocationPos::First
        } else if i + 1 == n {
            TreeLocationPos::Last
        } else {
            TreeLocationPos::Mid
        }
    }
}
