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

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);
}

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

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();
        
        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 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);

            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 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,
            });
            offset
        }
        Composite::Train2 { f, g } => {
            let f_offset = offset + 1;
            let f_pos = build_tree_from_expr(f, depth + 1, f_offset, nodes);

            let f_str = format!("{}", f);
            let g_offset = f_offset + f_str.chars().count() + 1;
            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,
            });
            f_pos
        }
        Composite::Train3 { f, g, h } => {
            let f_offset = offset + 1;
            let f_pos = build_tree_from_expr(f, depth + 1, f_offset, nodes);

            let f_str = format!("{}", f);
            let g_offset = f_offset + f_str.chars().count() + 1;
            let g_pos = build_tree_from_expr(g, depth + 1, g_offset, nodes);

            let g_str = format!("{}", g);
            let h_offset = g_offset + g_str.chars().count() + 1;
            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,
            });
            g_pos
        }
        Composite::Hook { kind, left, right } => {
            let left_pos = build_tree_from_expr(left, depth + 1, offset, nodes);

            let left_str = format!("{}", left);
            let kind_str = format!("{}", kind);
            let right_offset = offset + left_str.chars().count() + kind_str.chars().count();
            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,
            });
            left_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) => '┼',
    }
}