use std::collections::HashMap;
use std::cmp;
use std::sync::Arc;

use crate::{Env, Error, Exprs, FnArgs, FnReturn, Lines, Op, OpTrait, Pos, RefTVal, Token, TokenType, Tokens, Type, Value, VecMap, Precedence, Callable, umcore, unescape, parse_number};

/// The parsed representation of a script expression.
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum AST {
    /// A function call that accepts tokens as arguments.
    Macro {
        tokens: Vec<Token>,
    },
    /// A collection of other expressions.
    Container {
        token: Token,
        children: Vec<AST>,
    },
    /// An operation between two operands.
    Operator {
        token: Token,
        lhs: Box<AST>,
        rhs: Box<AST>,
    },
    /// A singular value.
    Value {
        token: Token,
    },
    /// The empty AST, used to represent an empty child. Cannot be run.
    Empty,
}
impl AST {
    pub(crate) fn from_value(v: &Value) -> Option<Self> {
        if let Value::Enum(_es, e) = v {
            if e.0 == "Empty" {
                return Some(AST::Empty);
            }

            let st = match &e.1.get(0) {
                Some(rv) => {
                    if let Value::Struct(st) = rv.clone_out().val {
                        st
                    } else {
                        return None
                    }
                },
                _ => return None,
            };
            match &*e.0 {
                "Macro" => return Some(AST::Macro {
                    tokens: match st.get("tokens") {
                        Some(rv) => {
                            if let Value::List(l) = rv.clone_out().val {
                                l.iter().map(|tv| Token::from_value(&tv.clone_out().val)).collect::<Option<Vec<Token>>>()?
                            } else {
                                return None
                            }
                        },
                        None => return None,
                    },
                }),
                "Container" => return Some(AST::Container {
                    token: match st.get("token") {
                        Some(tv) => Token::from_value(&tv.clone_out().val)?,
                        None => return None,
                    },
                    children: match st.get("children") {
                        Some(rv) => {
                            if let Value::List(l) = rv.clone_out().val {
                                l.iter().map(|tv| AST::from_value(&tv.clone_out().val)).collect::<Option<Vec<AST>>>()?
                            } else {
                                return None
                            }
                        },
                        None => return None,
                    },
                }),
                "Operator" => return Some(AST::Operator {
                    token: match st.get("token") {
                        Some(tv) => Token::from_value(&tv.clone_out().val)?,
                        None => return None,
                    },
                    lhs: match st.get("lhs") {
                        Some(tv) => Box::new(AST::from_value(&tv.clone_out().val)?),
                        None => return None,
                    },
                    rhs: match st.get("rhs") {
                        Some(tv) => Box::new(AST::from_value(&tv.clone_out().val)?),
                        None => return None,
                    },
                }),
                "Value" => return Some(AST::Value {
                    token: match st.get("token") {
                        Some(tv) => Token::from_value(&tv.clone_out().val)?,
                        None => return None,
                    },
                }),
                _ => {},
            }
        }
        None
    }

    pub(crate) fn get_first_token(&self) -> Option<Token> {
        match self {
            AST::Macro { tokens } => tokens.first().cloned(),
            AST::Container { token, .. } | AST::Value { token } => Some(token.clone()),
            AST::Operator { token, lhs, .. } => match **lhs {
                AST::Empty => Some(token.clone()),
                _ => lhs.get_first_token(),
            },
            AST::Empty => None,
        }
    }
    pub(crate) fn get_first_pos(&self) -> Option<Pos> {
        match self.get_first_token() {
            Some(t) => Some(t.pos),
            None => None,
        }
    }
    /**
    Returns a `String` representing the `AST`. Not guaranteed to be the same
    as the source script.
    */
    #[must_use]
    pub fn to_string_lossy(&self) -> String {
        match self {
            AST::Macro { tokens } => tokens.iter()
                .map(|t| t.data.clone())
                .collect::<Vec<String>>()
                .join(" "),
            AST::Container { token, children } => match &*token.data {
                "(" | "[" => {
                    let cs = children.iter()
                        .map(AST::to_string_lossy)
                        .collect::<Vec<String>>()
                        .join(", ");
                    format!("{}{}{}", token.data, cs, if token.data == "(" { ")" } else { "]" })
                },
                "{" => {
                    if children.is_empty() {
                        "{}".into()
                    } else {
                        let cs = children.iter()
                            .map(AST::to_string_lossy)
                            .collect::<Vec<String>>()
                            .join(";\n\t");
                        format!("{}\n\t{};\n{}", token.data, cs, "}")
                    }
                },
                s if s.starts_with("//") => token.data.clone(),
                _ => format!("{{{{Invalid container: {}}}}}", token.data),
            },
            AST::Operator { token, lhs, rhs } => match &*token.data {
                "(" | "[" | "{" => format!("{}{}", lhs.to_string_lossy(), rhs.to_string_lossy()),
                _ => format!("{}{}{}", lhs.to_string_lossy(), token.data, rhs.to_string_lossy()),
            },
            AST::Value { token } => token.data.clone(),
            AST::Empty => "EMPTY".into(),
        }
    }

    /// Parses the `Token` stream into a single `AST` representing the entire
    /// expression.
    ///
    /// # Errors
    /// Will return `Err` if `parse()` on children resulted in an error.
    /// An internal function will error if it's unable to reduce adjacent ASTs
    /// into a single AST.
    pub fn parse<I>(tokens: I, env: &Env<'_>) -> Result<AST, Error>
    where
        I: IntoIterator<Item = Token>
    {
        // Map tokens into flat ASTs
        let tokens: Vec<Token> = tokens.into_iter().collect();
        log::trace!("tokens: {:?}", tokens);

        let mut asts: Vec<Result<AST, Error>> = tokens.iter().enumerate()
            .fold(vec![], |mut acc: Vec<Result<AST, Error>>, (i, t)| {
                if let Some(Ok(AST::Macro { .. })) = acc.last() {
                    return acc;
                }

                match t.ttype {
                    TokenType::Identifier => {
                        if let Some(rv) = env.get(&t.data) {
                            if let Value::Function { args, .. } = rv.clone_out().val {
                                if args == umcore::_macro_args()
                                    || (args.contains_key("env") && args.contains_key("pos") && args.contains_key("tokens"))
                                {
                                    acc.push(Ok(AST::Macro { tokens: tokens[i..].to_vec() }));
                                    return acc;
                                }
                            }
                        }
                        acc.push(Ok(AST::Value { token: t.clone() }));
                        acc
                    },
                    TokenType::Number | TokenType::String => {
                        acc.push(Ok(AST::Value { token: t.clone() }));
                        acc
                    },
                    TokenType::Container => {
                        let children = Lines::split(&t.pos.filename, &t.data[1..cmp::max(1, t.data.len()-1)], Some(Pos {
                                filename: t.pos.filename.clone(),
                                line: t.pos.line,
                                col: t.pos.col+1,
                            })).flat_map(|l| Exprs::split(&l))
                            .map(|e| Tokens::tokenize(&e))
                            .map(|t| AST::parse(t, env))
                            .collect::<Result<Vec<AST>, Error>>();
                        match children {
                            Ok(cs) => acc.push(Ok(AST::Container {
                                token: Token {
                                    ttype: TokenType::Container,
                                    pos: t.pos.clone(),
                                    data: std::string::String::from(t.data.as_bytes()[0] as char),
                                },
                                children: cs,
                            })),
                            Err(m) => acc.push(Err(m)),
                        }
                        acc
                    },
                    TokenType::Symbol => {
                        acc.push(Ok(AST::Operator {
                            token: t.clone(),
                            lhs: Box::new(AST::Empty),
                            rhs: Box::new(AST::Empty),
                        }));
                        acc
                    },
                    TokenType::Comment => {
                        acc.push(Ok(AST::Container {
                            token: t.clone(),
                            children: vec![],
                        }));
                        acc
                    },
                }
            });
        for a in &asts {
            if let Err(m) = a {
                return Err(m.clone());
            }
        }
        let asts: Result<Vec<AST>, Error> = asts.drain(0..).collect();

        AST::group_precedence(asts?)
    }
    fn group_precedence(mut asts: Vec<AST>) -> Result<AST, Error> {
        let mut next_prec = 0;
        while asts.len() > 1 {
            let mut i: usize = 0;
            let mut unchanged = 0;
            while i <= asts.len() && asts.len() > 1 {
                if i == asts.len() {
                    i = 0;
                    unchanged += 1;
                    next_prec += 1;
                }
                // log::trace!("i {} unc {} prec {}", i, unchanged, next_prec);

                match asts[i].clone() {
                    // Reduce operators
                    AST::Operator { token, lhs, rhs } if *lhs == AST::Empty && *rhs == AST::Empty => {
                        if let Some(op) = Op::new(&token) {
                            if Op::prec(op) as usize <= next_prec {
                                log::trace!("Reducing op: {}", token);
                                // log::trace!("asts: {:?}", asts);
                                let lhs = if i > 0 {
                                    asts.remove(i-1)
                                } else {
                                    i += 1;
                                    AST::Empty
                                };
                                if i > 0 {
                                    let _op = asts.remove(i-1); // index changed from previous remove
                                } else {
                                    let _op = asts.remove(i);
                                    i += 1;
                                }
                                let rhs = if i > 0 && i <= asts.len() {
                                    asts.remove(i-1)
                                } else {
                                    AST::Empty
                                };

                                unchanged = 0;
                                asts.insert(i-1, AST::Operator {
                                    token,
                                    lhs: Box::new(lhs),
                                    rhs: Box::new(rhs),
                                });
                            }
                        }
                    },
                    // Reduce containers if they're for function calls, indexes, or constructors
                    AST::Container { token, children } if i > 0 && !token.data.starts_with("//") => {
                        if Precedence::L1 as usize <= next_prec {
                            log::trace!("Reducing container: {}", token);
                            // log::trace!("asts: {:?}", asts);
                            match &asts[i-1] {
                                AST::Operator { rhs, .. } if **rhs == AST::Empty =>  {
                                    // unchanged = 0;
                                },
                                _ => {
                                    let lhs = asts.remove(i-1);
                                    let _cont = asts.remove(i-1); // index changed from previous remove

                                    unchanged = 0;
                                    asts.insert(i-1, AST::Operator {
                                        token: token.clone(),
                                        lhs: Box::new(lhs),
                                        rhs: Box::new(AST::Container {
                                            token: Token {
                                                ttype: TokenType::Symbol, // Indicate function call with Symbol
                                                ..token
                                            },
                                            children,
                                        }),
                                    });
                                },
                            }
                        }
                    },

                    // Remove comments
                    AST::Container { token, .. } if token.data.starts_with("//") => {
                        asts.remove(i);
                        unchanged = 0;
                    },

                    AST::Macro { .. } | AST::Container { .. } => unchanged = 0,
                    // Reduce certain Values
                    AST::Value { token } => {
                        log::trace!("Reducing value: {}", token);
                        match token.ttype {
                            // Don't reduce sequential identifiers or strings
                            TokenType::Identifier | TokenType::String => {
                                if let Some(next) = asts.get(i+1) {
                                    match next {
                                        AST::Macro { .. } => unchanged += 1,
                                        AST::Value { token } => match token.ttype {
                                            TokenType::Identifier | TokenType::Number | TokenType::String => unchanged += 1,
                                            _ => unchanged = 0,
                                        },
                                        _ => unchanged = 0,
                                    }
                                } else {
                                    unchanged = 0;
                                }
                            },
                            // Combine number tokens
                            TokenType::Number => {
                                if let Some(next) = asts.get(i+1).cloned() {
                                    match next {
                                        AST::Value { token: t2 } => {
                                            asts.remove(i+1);
                                            asts.remove(i);

                                            asts.insert(i, AST::Value {
                                                token: Token {
                                                    ttype: TokenType::Number,
                                                    pos: token.pos,
                                                    data: token.data + &t2.data,
                                                },
                                            });

                                            unchanged = 0;
                                        },
                                        AST::Container { token, .. } if token.data.starts_with("//") => unchanged = 0,
                                        AST::Operator { .. } => unchanged = 0,
                                        _ => unchanged += 1,
                                    }
                                }
                            },
                            _ => unchanged = 0,
                        }
                    },
                    _ => {
                        log::trace!("{:?}", asts[i]);
                    },
                }
                // Error out if unable to reduce
                if unchanged > 0 {
                    if next_prec < Precedence::MAX as usize {
                        i = 0;
                        unchanged = 0;
                        next_prec += 1;
                        continue;
                    }

                    return Err(Error::Parse(
                        format!("Unreducible values [{}/{}]: {:?}", i+1, asts.len(), asts.iter()
                            .map(|a| "<".to_string()+&a.to_string_lossy()+">")
                            .intersperse(" ".into())
                            .collect::<String>()),
                        asts[0].get_first_pos())
                    );
                }
                i += 1;
            }
        }
        if asts.is_empty() {
            Ok(AST::Empty)
        } else {
            Ok(asts.remove(0))
        }
    }

    /**
    Runs the `AST` in the given environment and returns the value of the
    expression.
    */
    pub fn run(&self, env: &Env<'_>) -> FnReturn {
        match self {
            AST::Macro { tokens } => {
                let t0 = tokens.first().expect("expected macro token");
                if let Some(rv) = env.get(&t0.data) {
                    if let Value::Function { body, .. } = rv.clone_out().val {
                        return body.call(&*env, FnArgs::Macro {
                            vars: Env::flatten(env).vars.into(),
                            pos: self.get_first_pos(),
                            tokens: tokens.clone(),
                        });
                    }
                }
                return (None, Err(Error::Script(format!("failed to find macro {} in env", t0.data), self.get_first_pos())));
            },
            AST::Container { token, children } => {
                let mut nenv = Env::child(env);
                let mut should_map = false;
                let vals: Result<Vec<RefTVal>, Error> = children.iter()
                    .filter_map(|ast| {
                        if let AST::Operator { token, .. } = &ast {
                            if token.data == ":" {
                                should_map = true;
                            }
                        } else if let AST::Container { token, .. } = &ast {
                            // Skip comments
                            if token.data.starts_with("//") {
                                return None;
                            }
                        }

                        // Run child AST
                        let (vars, v) = ast.run(&nenv);
                        if let Some(vars) = vars {
                            log::trace!("{:?}", vars.keys());
                            nenv.update(vars);
                            // crate::run_interactive(&nenv);
                        }
                        Some(v)
                    }).collect();
                let diff = env.diff(nenv.clone());
                let vals: Vec<RefTVal> = match vals {
                    Ok(vs) => vs,
                    Err(m) => return (Some(diff), Err(m)),
                };
                match &*token.data {
                    "(" => {
                        // Return a list if it's for a function call
                        // if vals.len() > 1 || vals.is_empty() {
                        if token.ttype == TokenType::Symbol {
                            return (Some(diff), Ok(Value::List(vals).into()))
                        }
                        return (Some(diff), Ok(vals.last().unwrap_or(&Value::none().into()).clone()))
                    },
                    "[" => return (Some(nenv.vars), Ok(Value::List(vals).into())),
                    "{" => {
                        // Attempt to build Map from kv pairs
                        if should_map {
                            let map: Result<Vec<(RefTVal, RefTVal)>, Error> = vals.iter()
                                .map(|rv| {
                                    if let Value::List(l) = rv.clone_out().val {
                                        if l.len() == 2 {
                                            return Ok((l[0].clone(), l[1].clone()));
                                        }
                                    }
                                    return Err(Error::Script(format!("failed to build Map from Container{{}}: {} is not a List", rv), Some(token.pos.clone())));
                                }).collect();
                            if let Ok(m) = map {
                                return (Some(diff), Ok(Value::Map(m.into_iter().collect::<HashMap<RefTVal, RefTVal>>().into()).into()));
                            }
                        }
                        return (Some(diff), Ok(vals.last().unwrap_or(&Value::none().into()).clone()));
                    },
                    _ => {
                        return (None, Err(Error::Script(format!("unsupported Container {}", token.data), self.get_first_pos())));
                    },
                }
            },
            AST::Operator { token, lhs, rhs } => {
                let mut lval = None;
                let mut rval = None;

                // Handle the Dollar operator without evaluating either side
                if token.data == "$" {
                    let args = match &**lhs {
                        AST::Value { token: t2 } => Arc::new(vec![
                            (t2.data.clone(), Type::any()),
                        ]),
                        _ => return (None, Err(Error::Script(format!("invalid function parameter: {}", lhs.to_string_lossy()), Some(token.pos.clone())))),
                    };
                    return (None, Ok(Value::Function {
                        args,
                        vars: Env::flatten(env).vars.into(),
                        body: Callable::AST(*rhs.clone()),
                    }.into()));
                }

                // Handle the Equal operator without evaluating the LHS
                if token.data == "=" {
                    let lval: RefTVal = match &**lhs {
                        AST::Value { token: lt } => Value::String(lt.data.clone()).into(),
                        AST::Operator { token: op, .. } if op.data == "." || op.data == "@" => {
                            match lhs.run(env).1 {
                                Ok(v) => v,
                                Err(m) => return (None, Err(m)),
                            }
                        },
                        AST::Operator { token: op, lhs: llhs, rhs: lrhs } if op.data == "*" && **llhs == AST::Empty => {
                            match lrhs.run(env).1 {
                                Ok(v) => v,
                                Err(m) => return (None, Err(m)),
                            }
                        },
                        _ => return (None, Err(Error::Script(format!("invalid LHS for Equal: {:?}", lhs), Some(token.pos.clone())))),
                    };
                    let rval = match rhs.run(env).1 {
                        Ok(v) => Some(v),
                        Err(m) => return (None, Err(m)),
                    };
                    return OpTrait::<{Op::Equal}>::op(&lval, env, token.pos.clone(), rval.as_ref());
                }

                // Compute lval
                if **lhs != AST::Empty {
                    lval = match lhs.run(env).1 {
                        Ok(v) => Some(v),
                        Err(m) => return (None, Err(m)),
                    };
                }

                // Handle Dot and Ampersat operators without evaluating the RHS
                if token.data == "." || token.data == "@" {
                    let (rstr, rval): (String, Option<RefTVal>) = match &**rhs {
                        AST::Value { token: rtoken } => {
                            let rstr = rtoken.data.clone();
                            (rstr.clone(), Some(Value::String(rstr).into()))
                        },
                        AST::Container { .. } => {
                            // Only evaluate RHS if it's in a container
                            match rhs.run(env) {
                                (_, Ok(v)) => ("".into(), Some(v)),
                                (vars, Err(m)) => return (vars, Err(m)),
                            }
                        },
                        _ => {
                            let rstr = match rhs.get_first_token() {
                                Some(t) => t.data,
                                None => return (None, Err(Error::Script(format!("expected RHS token for operator {}", token.data), Some(token.pos.clone())))),
                            };
                            (rstr.clone(), Some(Value::String(rstr).into()))
                        },
                    };
                    match lval {
                        Some(lval) => return Op::op(env, token.clone(), Some(&lval), rval.as_ref()),
                        None if token.data == "." => {
                            // If no LHS then fetch from the Enums in the current env
                            for (_pek, pev) in Env::iter_flatten(env) {
                                if let Ok(pev) = pev.lock() {
                                    if let Value::Type(en, enr) = &pev.val {
                                        if let Type::Enum(envar) = &**en {
                                            if enr.get(&rstr).is_some() {
                                                let v = Value::Enum(Arc::clone(envar), Box::new((rstr, vec![])));
                                                return (None, Ok(v.into()));
                                            }
                                        }
                                    }
                                }
                            }
                        },
                        None if token.data == "@" => {},
                        None => {},
                    }
                }

                // Compute rval
                if **rhs != AST::Empty {
                    rval = match rhs.run(env).1 {
                        Ok(v) => Some(v),
                        Err(m) => return (None, Err(m)),
                    };
                }

                // Compute operation
                let (vars, val) = Op::op(env, token.clone(), lval.as_ref(), rval.as_ref());
                let nvars = vars.map(|vs| env.diff(Env::from(vs)));
                return (nvars, val);
            },
            AST::Value { token } => {
                match token.ttype {
                    TokenType::Identifier => match env.get(&token.data) {
                        Some(v) => return (None, Ok(v.clone())),
                        None => return (None, Err(Error::Script(format!("failed to find '{}' in this env", token.data), self.get_first_pos()))),
                    },
                    TokenType::Number => return parse_number(&token.data, self.get_first_pos()),
                    TokenType::String => return (None, Ok(Value::String(unescape(&token.data[1..token.data.len()-1])).into())),
                    _ => {},
                }
            },
            AST::Empty => {},
        }
        (None, Err(Error::Script(format!("failed to execute AST: {:?}", self), self.get_first_pos())))
    }
}