use rprt_engine::{expression::*, parser::*, token::*};
use std::collections::BTreeSet;
use std::env;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: {} <code>", args[0]);
        std::process::exit(1);
    }

    let input = &args[1];

    let tokens = match tokenise(input) {
        Ok(t) => t,
        Err(e) => {
            eprintln!("Tokenization error: {}", e);
            std::process::exit(1);
        }
    };

    let parsed = match parse(tokens) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("Parse error: {:?}", e);
            std::process::exit(1);
        }
    };

    let formatted = format!("{}", parsed);
    let tree = generate_tree(&formatted, &parsed);
    print!("{}", tree);
}

enum NodeKind {
    Hook(HookKind),
    Train(u32),
    Group,
    Plain,
}

struct TreeNode {
    pos: usize,
    depth: usize,
    children: Vec<usize>,
    is_leaf: bool,
    node_kind: NodeKind,
}

fn generate_tree(formatted: &str, expr: &Composite) -> String {
    let width = formatted.chars().count();

    let mut nodes = Vec::new();
    build_tree_from_expr(expr, 0, 0, &mut nodes);

    let max_depth = nodes.iter().map(|n| n.depth).max().unwrap_or(0);

    let mut output = String::new();
    output.push_str(formatted);
    output.push('\n');

    let mut active_bars: BTreeSet<usize> =
        nodes.iter().filter(|n| n.is_leaf).map(|n| n.pos).collect();

    let mut prev_active = active_bars.clone();

    for current_depth in (0..=max_depth).rev() {
        let interior: Vec<_> = nodes
            .iter()
            .filter(|n| n.depth == current_depth && !n.is_leaf)
            .collect();

        if interior.is_empty() {
            continue;
        }

        let mut next_active = active_bars.clone();
        for node in &interior {
            for &child in &node.children {
                next_active.remove(&child);
            }
            next_active.insert(node.pos);
        }

        let mut horizontal_spans = Vec::new();
        for node in &interior {
            if !node.children.is_empty() {
                let left = *node.children.iter().min().unwrap();
                let right = *node.children.iter().max().unwrap();
                horizontal_spans.push((left, right));
            }
        }

        let mut spacer = vec![' '; width];
        for i in 0..width {
            if prev_active.contains(&i) {
                spacer[i] = '│';
            }
        }
        output.push_str(&spacer.iter().collect::<String>());
        output.push('\n');

        let mut line = vec![' '; width];
        for i in 0..width {
            let up = prev_active.contains(&i);
            let down = next_active.contains(&i);
            let in_span = horizontal_spans.iter().any(|(l, r)| *l <= i && i <= *r);
            let left = i > 0
                && horizontal_spans
                    .iter()
                    .any(|(l, r)| *l <= i - 1 && i - 1 <= *r);
            let right = i + 1 < width
                && horizontal_spans
                    .iter()
                    .any(|(l, r)| *l <= i + 1 && i + 1 <= *r);

            let node_at_pos = interior.iter().find(|n| n.pos == i);

            if let Some(node) = node_at_pos {
                match node.node_kind {
                    NodeKind::Hook(hook_kind) => {
                        line[i] = hook_kind.to_string().chars().next().unwrap()
                    }
                    NodeKind::Train(arity) => line[i] = char::from_digit(arity as u32, 10).unwrap(),
                    NodeKind::Group => line[i] = '+',
                    NodeKind::Plain => {
                        if up || down || in_span {
                            line[i] = box_char(up, down, left, right);
                        }
                    }
                }
            } else if up || down || in_span {
                line[i] = box_char(up, down, left, right);
            }
        }

        output.push_str(&line.iter().collect::<String>());
        output.push('\n');

        prev_active = next_active.clone();
        active_bars = next_active;
    }

    if !nodes.is_empty() {
        let mut final_line = vec![' '; width];
        if let Some(&pos) = prev_active.iter().next() {
            if pos < width {
                final_line[pos] = '│';
            }
        }
        output.push_str(&final_line.iter().collect::<String>());
        output.push('\n');
    }

    output
}

fn format_width(expr: &Composite) -> usize {
    format!("{}", expr).chars().count()
}

fn build_tree_from_expr(
    expr: &Composite,
    depth: usize,
    offset: usize,
    nodes: &mut Vec<TreeNode>,
) -> usize {
    match expr {
        Composite::SelectionFunction { .. } | Composite::TextFunction { .. } => {
            nodes.push(TreeNode {
                pos: offset,
                depth,
                children: vec![],
                is_leaf: true,
                node_kind: NodeKind::Plain,
            });
            offset
        }
        Composite::Train2 { f, g } => {
            let f_offset = offset + 1; // Account for opening '('
            let f_pos = build_tree_from_expr(f, depth + 1, f_offset, nodes);
            let g_offset = f_offset + format_width(f) + 1; // +1 for space
            let g_pos = build_tree_from_expr(g, depth + 1, g_offset, nodes);

            nodes.push(TreeNode {
                pos: f_pos,
                depth,
                children: vec![f_pos, g_pos],
                is_leaf: false,
                node_kind: NodeKind::Train(2),
            });
            f_pos
        }
        Composite::Train3 { f, g, h } => {
            let f_offset = offset + 1; // Account for opening '('
            let f_pos = build_tree_from_expr(f, depth + 1, f_offset, nodes);
            let g_offset = f_offset + format_width(f) + 1; // +1 for space
            let g_pos = build_tree_from_expr(g, depth + 1, g_offset, nodes);
            let h_offset = g_offset + format_width(g) + 1; // +1 for space
            let h_pos = build_tree_from_expr(h, depth + 1, h_offset, nodes);

            nodes.push(TreeNode {
                pos: g_pos,
                depth,
                children: vec![f_pos, g_pos, h_pos],
                is_leaf: false,
                node_kind: NodeKind::Train(3),
            });
            g_pos
        }
        Composite::Hook { kind, left, right } => {
            let left_offset = offset + 1;
            let left_pos = build_tree_from_expr(left, depth + 1, left_offset, nodes);
            let right_offset = left_offset + format_width(left) + 1;
            let right_pos = build_tree_from_expr(right, depth + 1, right_offset, nodes);

            nodes.push(TreeNode {
                pos: left_pos,
                depth,
                children: vec![left_pos, right_pos],
                is_leaf: false,
                node_kind: NodeKind::Hook(*kind),
            });
            left_pos
        }
        Composite::Group { operations } => {
            let mut children = Vec::new();
            let mut current_offset = offset + 2;

            for (i, op) in operations.iter().enumerate() {
                if i > 0 {
                    current_offset += 3;
                }
                let pos = build_tree_from_expr(op, depth + 1, current_offset, nodes);
                children.push(pos);
                current_offset += format_width(op);
            }

            let node_pos = children.first().copied().unwrap_or(offset);
            nodes.push(TreeNode {
                pos: node_pos,
                depth,
                children,
                is_leaf: false,
                node_kind: NodeKind::Group,
            });
            node_pos
        }
    }
}

fn box_char(up: bool, down: bool, left: bool, right: bool) -> char {
    match (up, down, left, right) {
        (false, false, false, false) => ' ',
        (true, false, false, false) => '│',
        (false, true, false, false) => '│',
        (true, true, false, false) => '│',
        (false, false, true, false) => '─',
        (false, false, false, true) => '─',
        (false, false, true, true) => '─',
        (true, false, true, false) => '┘',
        (true, false, false, true) => '└',
        (false, true, true, false) => '┐',
        (false, true, false, true) => '┌',
        (true, false, true, true) => '┴',
        (false, true, true, true) => '┬',
        (true, true, true, false) => '┤',
        (true, true, false, true) => '├',
        (true, true, true, true) => '┼',
    }
}