use runtime::*;
use parser::*;
use std::collections::HashMap;
use std::rc::Rc;

struct ElementPartsCollector {
    res: String, cur_col: usize,
}

impl ElementPartsCollector {
    fn without_pad(&mut self, thp: &str) {
        self.res += thp;
        for ch in thp.chars() {
            if ch == '\n' { self.cur_col = 0; } else { self.cur_col += 1; }
        }
    }

    fn add_padded_element_piece(&mut self, offset: usize, thp: &str){
        for ch in thp.chars() {
            self.res.push(ch);
            if ch == '\n' {
                self.res.push_str(&" ".repeat(offset));
                self.cur_col = offset;
            } else {
                self.cur_col += 1;
            }
        }
    }

    fn add_single_padded_element_piece(&mut self, thp: &str) {
        self.add_padded_element_piece(self.cur_col, thp);
    }
}

impl<'m> RuntimeEnv<'m> {
    fn get_needed_if_block<'i>(
        &self, if_sub_el: &'i IfSubElement, local: &[Value<'m>], recc: u32
    ) -> Result<Option<&'i Element>, String> {
        let n = if_sub_el.conditions.len();
        for i in 0..n {
            let r = self.execute_expression(&if_sub_el.conditions[i], local, recc - 1)?;
            if self.value_to_bool(&r)? {
                return Ok(Some(&if_sub_el.branches[i]));
            }
        }
        if if_sub_el.branches.len() > n {
            Ok(Some(&if_sub_el.branches[n]))
        } else {
            Ok(None)
        }
    }

    /* if Ok, local remains unchanged after return, but if Err, don't use local anymore  */
    fn execute_element_block(
        &self, el: &[SubElement], local: &mut Vec<Value<'m>>, recc: u32
    ) -> Result<String, String> {
        if recc == 0 { return Err(self.make_error("Recursion limit exceeded")) }
        let mut res = ElementPartsCollector{res: String::new(), cur_col: 0};
        for se in el {
            match se {
                SubElement::Static(ps) => {
                    res.without_pad(ps);
                },
                SubElement::If(if_sub_el) => if let Some(branch_el) =
                    self.get_needed_if_block(if_sub_el, local, recc - 1)? {
                        res.add_single_padded_element_piece(
                            &self.execute_element_block(branch_el, local, recc - 1)?)
                },
                SubElement::InsertExpr(expr) => {
                    let val = self.execute_expression(expr, local, recc - 1)?;
                    match val.if_str_get_ref() {
                        Some(s) => res.add_single_padded_element_piece(s),
                        None => return Err(self.make_error(&format!(
                            "Can't insert {} into template. Consider using {{{{}}}}", val.type_name()
                        )))
                    }
                },
                SubElement::For(ForSubElement{iterable, core, separator}) => {
                    let _g = self.register("In {%for%}".into());
                    let r = self.execute_expression(iterable, local, recc - 1)?;
                    let saved_offset = res.cur_col;
                    let k = local.len();
                    local.push(Value::Int(0)); local.push(Value::Int(0));
                    let collected: Vec<String> = match r {
                        Value::Dict(dict) => dict.as_ref().iter()
                            .map(|(key ,value)| -> Result<String, String> {
                                local[k] = Value::Str(Rc::new(key.clone()));
                                local[k + 1] = value.clone();
                                self.execute_element_block(core, local, recc - 1)
                            }).collect::<Result<Vec<String>, String>>()?,
                        Value::Ref(SharedValue::Dict(dict)) => dict.iter()
                            .map(|(key, sh_val)| -> Result<String, String> {
                                local[k] = Value::Str(Rc::new(key.clone()));
                                local[k + 1] = Value::Ref(sh_val);
                                self.execute_element_block(core, local, recc - 1)
                            }).collect::<Result<Vec<String>, String>>()?,
                        Value::Arr(arr) => arr.as_ref().iter().enumerate()
                            .map(|(ind ,value)| -> Result<String, String> {
                                local[k] = Value::Int(ind as i64);
                                local[k + 1] = value.clone();
                                self.execute_element_block(core, local, recc - 1)
                            }).collect::<Result<Vec<String>, String>>()?,
                        Value::Ref(SharedValue::Arr(arr)) => arr.iter().enumerate()
                            .map(|(ind, sh_val)| -> Result<String, String> {
                                local[k] = Value::Int(ind as i64);
                                local[k + 1] = Value::Ref(sh_val);
                                self.execute_element_block(core, local, recc - 1)
                            }).collect::<Result<Vec<String>, String>>()?,
                        _ => return Err(self.make_error(&format!("Can't iterate over {}", r.type_name())))
                    };
                    if !collected.is_empty() {
                        res.add_padded_element_piece(saved_offset, &collected[0]);
                        for s in &collected[1..]{
                            res.add_padded_element_piece(saved_offset, separator);
                            res.add_padded_element_piece(saved_offset, s);
                        }
                    }
                }
                SubElement::Let(expr, core) => {
                    let r = self.execute_expression(expr, local, recc - 1)?;
                    local.push(r);
                    res.add_single_padded_element_piece(&self.execute_element_block(
                        &core, local, recc - 1)?);
                }
            };
        }
        Ok(res.res)
    }
}

/* We construct intermediate closures for elements at runtime
 * argc >= 1 */
impl<'m> RuntimeEnv<'m> {
    fn construct_result_of_element_lambda(
        &self, mut prev: Vec<Value<'m>>, argc: usize, el_source_pos: usize, recc: u32,
    ) -> Result<Value<'m>, String> {
        if prev.len() == argc {
            self.execute_element_block(&self.mtgott.source_elements[el_source_pos], &mut prev, recc)
                .map(|s| Value::Str(Rc::new(s)))
        } else {
            Ok(Value::Fn(Rc::new(move |re, arg, counter| -> Result<Value<'m>, String> {
                re.check_recc(counter)?;
                let mut prev = prev.clone();
                prev.push(arg);
                re.construct_result_of_element_lambda(prev, argc, el_source_pos, counter - 1)
            } )))
        }
    }
}

/* This function is supposed to compile all elements in one particular file.
 Due to how bad mtgott type system is, some elements require me to save source tree into
 a storage */
pub fn plemege_to_value(
    x: Plemege, file_path: &str,
    source_elements: &mut Vec<Element>, source_expr: &mut Vec<Expression>, recc: u32
) -> Result<SharedValue, String> {
    if recc == 0 { return Err("Recursion limit exceeded".into()) }
    match x {
        Plemege::Package(map) => {
            let mut new_dict: HashMap<String, SharedValue> = HashMap::new();
            for (key, thing) in map {
                new_dict.insert(key, plemege_to_value(thing, file_path, source_elements, source_expr, recc - 1)?);
            }
            Ok(SharedValue::Dict(new_dict))
        },
        Plemege::Element(argc, el) => {
            if argc == 0 && el.iter().all(|se| matches!(se, SubElement::Static(_))) {
                /* The only optimization we ever make (we could do more, but sorry, my life is finite) */
                Ok(SharedValue::Str(el.into_iter()
                    .fold(String::new(), |mut acc, p| {
                    let SubElement::Static(s) = p else { panic!() };
                        acc.push_str(&s); acc
                })))
            } else if argc == 0 {
                Ok(SharedValue::RuntimeConst(Box::new(
                    move |re: &RuntimeEnv, counter: u32| -> Result<Value, String> {
                        re.check_recc(counter)?;
                        /* No need to save element into source vector. We move it here */
                        re.execute_element_block(&el, &mut Vec::new(), counter).map(|s| Value::Str(Rc::new(s)))
                    })))
            } else {
                let body_element_id = source_elements.len();
                source_elements.push(el);
                Ok(SharedValue::Fn(Box::new(
                    move |re, x, counter| -> Result<Value, String> {
                        re.check_recc(counter)?;
                        re.construct_result_of_element_lambda(vec![x], argc, body_element_id, counter - 1)
                    }
                )))
            }
        }
        Plemege::Expression(expr) => match expr {
            Expression::Int(n) => Ok(SharedValue::Int(n as i64)),
            Expression::Str(s) => Ok(SharedValue::Str(s)),
            // todo: optimize arrays
            Expression::Lambda(NewLambdaExpression{local_var_array, expr_id}) => {
                assert_eq!(local_var_array, vec![usize::MAX]);
                Ok(SharedValue::Fn(Box::new(
                    move |re, arg, counter| -> Result<Value, String> {
                        re.check_recc(counter)?;
                        re.execute_expression(
                            &re.mtgott.source_expressions[expr_id],
                            &vec![arg], counter - 1)
                    }))
                )
            },
            _ => Ok(SharedValue::RuntimeConst(Box::new(
                move |re, counter| -> Result<Value, String> {
                    re.check_recc(counter)?;
                    /* No need to save expression into source vector. We move it here */
                    re.execute_expression(&expr, &Vec::new(), counter - 1)
                })))
        }
    }
}