use std::collections::{HashSet, HashMap};
use std::iter::FromIterator;
use vm::representations::{PreSymbolicExpression, ClarityName};
use vm::representations::PreSymbolicExpressionType::{AtomValue, Atom, List, Tuple, SugaredContractIdentifier, SugaredFieldIdentifier, TraitReference, FieldIdentifier};
use vm::functions::NativeFunctions;
use vm::functions::define::DefineFunctions;
use vm::ast::types::{ContractAST, BuildASTPass};
use vm::ast::errors::{ParseResult, ParseError, ParseErrors};
use vm::costs::{CostTracker, cost_functions};
use vm::types::{Value};

#[cfg(test)]
mod tests;

pub struct DefinitionSorter {
    graph: Graph,
    top_level_expressions_map: HashMap<ClarityName, TopLevelExpressionIndex>   
}

impl <'a> DefinitionSorter {

    fn new() -> Self {
        Self { 
            top_level_expressions_map: HashMap::new(),
            graph: Graph::new()
        }
    }

    pub fn run_pass<T: CostTracker>(contract_ast: &mut ContractAST, accounting: &mut T) -> ParseResult<()> {
        let mut pass = DefinitionSorter::new();
        pass.run(contract_ast, accounting)?;
        Ok(())
    }

    pub fn run<T: CostTracker>(&mut self, contract_ast: &mut ContractAST, accounting: &mut T) -> ParseResult<()> {
        let exprs = contract_ast.pre_expressions[..].to_vec();
        for (expr_index, expr) in exprs.iter().enumerate() {
            self.graph.add_node(expr_index);

            match self.find_expression_definition(expr) {
                Some((definition_name, atom_index, _)) => {
                    let tle = TopLevelExpressionIndex { expr_index, atom_index };
                    self.top_level_expressions_map.insert(definition_name, tle);
                },
                None => {}
            }
        }

        for (expr_index, expr) in exprs.iter().enumerate() {
            self.probe_for_dependencies(&expr, expr_index)?;
        }

        runtime_cost!(cost_functions::AST_CYCLE_DETECTION, accounting, self.graph.edges_count()?)?;

        let mut walker = GraphWalker::new();
        let sorted_indexes = walker.get_sorted_dependencies(&self.graph)?;
        
        if let Some(deps) = walker.get_cycling_dependencies(&self.graph, &sorted_indexes) {
            let mut deps_props = vec![];
            for i in deps.iter() {
                let exp = &contract_ast.pre_expressions[*i];
                if let Some(def) = self.find_expression_definition(&exp) {
                    deps_props.push(def);
                }
            }
            let functions_names = deps_props.iter().map(|i| i.0.to_string()).collect();

            let error = ParseError::new(ParseErrors::CircularReference(functions_names));
            return Err(error)
        }

        contract_ast.top_level_expression_sorting = Some(sorted_indexes);
        Ok(())
    }

    fn probe_for_dependencies(&mut self, expr: &PreSymbolicExpression, tle_index: usize) -> ParseResult<()> {
        match expr.pre_expr {
            Atom(ref name) => {
                if let Some(dep) = self.top_level_expressions_map.get(name) {
                    if dep.atom_index != expr.id {
                        self.graph.add_directed_edge(tle_index, dep.expr_index);
                    }
                }
                Ok(())
            },
            TraitReference(ref name) => { 
                if let Some(dep) = self.top_level_expressions_map.get(name) {
                    if dep.atom_index != expr.id {
                        self.graph.add_directed_edge(tle_index, dep.expr_index);
                    }
                }
                Ok(())
            },
            List(ref exprs) => {
                // Avoid looking for dependencies in tuples
                // TODO: Eliminate special handling of tuples as it is a separate presymbolic expression type
                if let Some((function_name, function_args)) = exprs.split_first() {
                    if let Some(function_name) = function_name.match_atom() {
                        if let Some(define_function) = DefineFunctions::lookup_by_name(function_name) {
                            match define_function {
                                DefineFunctions::PersistedVariable | DefineFunctions::Constant  => {
                                    // Args: [(define-name-and-types), ...]: ignore 1st arg
                                    if function_args.len() > 1 {
                                        for expr in function_args[1..function_args.len()].into_iter() {
                                            self.probe_for_dependencies(expr, tle_index)?;
                                        }
                                    }
                                    return Ok(());
                                },
                                DefineFunctions::PublicFunction | DefineFunctions::PrivateFunction |
                                DefineFunctions::ReadOnlyFunction => {
                                    // Args: [(define-name-and-types), ...]
                                    if function_args.len() == 2 {
                                        self.probe_for_dependencies_in_define_args(&function_args[0], tle_index)?;
                                        self.probe_for_dependencies(&function_args[1], tle_index)?;
                                    } 
                                    return Ok(());
                                },
                                DefineFunctions::Map => {
                                    // Args: [name, tuple-key, tuple-value]: handle tuple-key and tuple-value as tuples
                                    if function_args.len() == 3 {
                                        self.probe_for_dependencies_in_tuple(&function_args[1], tle_index)?;
                                        self.probe_for_dependencies_in_tuple(&function_args[2], tle_index)?;
                                    }
                                    return Ok(());
                                },
                                DefineFunctions::Trait => {
                                    if function_args.len() != 2 {
                                        return Ok(())
                                    } 
                                    if let Some(trait_sig) = function_args[1].match_list() {
                                        for func_sig in trait_sig.iter() {
                                            if let Some(func_sig) = func_sig.match_list() {
                                                if func_sig.len() == 3 {
                                                    self.probe_for_dependencies(&func_sig[1], tle_index)?;
                                                    self.probe_for_dependencies(&func_sig[2], tle_index)?;
                                                }     
                                            }
                                        }
                                    }
                                    return Ok(())
                                },
                                DefineFunctions::ImplTrait | DefineFunctions::UseTrait => {
                                    return Ok(())
                                },
                                DefineFunctions::NonFungibleToken | DefineFunctions::FungibleToken => {
                                    return Ok(())
                                }
                            }
                        } else if let Some(native_function) = NativeFunctions::lookup_by_name(function_name) {
                            match native_function {
                                NativeFunctions::FetchEntry | NativeFunctions::DeleteEntry => {
                                    // Args: [map-name, tuple-predicate]: handle tuple-predicate as tuple
                                    if function_args.len() == 2 {
                                        self.probe_for_dependencies(&function_args[0], tle_index)?;
                                        self.probe_for_dependencies_in_tuple(&function_args[1], tle_index)?;
                                    }
                                    return Ok(());
                                }, 
                                NativeFunctions::SetEntry | NativeFunctions::InsertEntry => {
                                    // Args: [map-name, tuple-keys, tuple-values]: handle tuple-keys and tuple-values as tuples
                                    if function_args.len() == 3 {
                                        self.probe_for_dependencies(&function_args[0], tle_index)?;
                                        self.probe_for_dependencies_in_tuple(&function_args[1], tle_index)?;
                                        self.probe_for_dependencies_in_tuple(&function_args[2], tle_index)?;
                                    }
                                    return Ok(());
                                }, 
                                NativeFunctions::ContractCall => {
                                    // Args: [contract-name, function-name, ...]: ignore contract-name, function-name, handle rest
                                    if function_args.len() > 2 {
                                        for expr in function_args[2..].iter() {
                                            self.probe_for_dependencies(expr, tle_index)?;
                                        }
                                    }
                                    return Ok(());
                                },
                                NativeFunctions::Let => {
                                    // Args: [((name-1 value-1) (name-2 value-2)), ...]: handle 1st arg as a tuple
                                    if function_args.len() > 1 {
                                        self.probe_for_dependencies_in_tuple(&function_args[0], tle_index)?;
                                        for expr in function_args[1..function_args.len()].into_iter() {
                                            self.probe_for_dependencies(expr, tle_index)?;
                                        }
                                    }
                                    return Ok(());
                                }
                                NativeFunctions::Len => {
                                    // Args: buffer | list
                                    if function_args.len() == 1 {
                                        self.probe_for_dependencies(&function_args[0], tle_index)?;
                                    }
                                    return Ok(());
                                },
                                NativeFunctions::TupleGet => {
                                    // Args: [key-name, expr]: ignore key-name
                                    if function_args.len() == 2 {
                                        self.probe_for_dependencies(&function_args[1], tle_index)?;
                                    }
                                    return Ok(());
                                },
                                NativeFunctions::TupleCons => {
                                    // Args: [(key-name A), (key-name-2 B), ...]: handle as a tuple
                                    self.probe_for_dependencies_in_tuple_list(function_args, tle_index)?;
                                    return Ok(());
                                },
                                _ => {}
                            }
                        }
                    }
                }
                for expr in exprs.into_iter() {
                    self.probe_for_dependencies(expr, tle_index)?;
                }
                Ok(())
            },
            Tuple(ref exprs) => {
                self.probe_for_dependencies_in_tuple_list(exprs, tle_index)?;
                Ok(())
            },
            AtomValue(_) | FieldIdentifier(_) | SugaredContractIdentifier(_) | SugaredFieldIdentifier(_, _) => Ok(()),
        }
    }

    /// accept a slice of expected-pairs, e.g., [ (a b) (c d) (e f) ], and
    ///   probe them for dependencies as if they were part of a tuple definition.
    fn probe_for_dependencies_in_tuple_list(&mut self, pairs: &[PreSymbolicExpression], tle_index: usize) -> ParseResult<()> {
        for pair in pairs.iter() {
            if let Some(pair) = pair.match_list() {
                if pair.len() == 2 {
                    self.probe_for_dependencies(&pair[1], tle_index)?;
                }
            }
        }
        Ok(())
    }
    
    fn probe_for_dependencies_in_define_args(&mut self, expr: &PreSymbolicExpression, tle_index: usize) -> ParseResult<()> {
        if let Some(func_sig) = expr.match_list() {
            // Func definitions can look like:
            // 1. (define-public func_name body)
            // 2. (define-public (func_name (arg uint) ...) body)
            // The goal here is to traverse case 2, looking for trait references
            if let Some((_, args)) = func_sig.split_first() {
                for pair in args.into_iter() {
                    if let Some(pair) = pair.match_list() {
                        if pair.len() == 2 {
                            self.probe_for_dependencies(&pair[1], tle_index)?;
                        }
                    }
                }    
            }
        }
        Ok(())
    }

    fn probe_for_dependencies_in_tuple(&mut self, expr: &PreSymbolicExpression, tle_index: usize) -> ParseResult<()> {
        if let Some(tuple) = expr.match_list() {
            self.probe_for_dependencies_in_tuple_list(tuple, tle_index)?;
        }
        Ok(())
    }

    fn find_expression_definition<'b>(&mut self, exp: &'b PreSymbolicExpression) -> Option<(ClarityName, u64, &'b PreSymbolicExpression)> {
        let args = {
            let exp = exp.match_list()?;
            let (function_name, args) = exp.split_first()?;
            let function_name = function_name.match_atom()?;
            DefineFunctions::lookup_by_name(function_name)?;
            Some(args)
        }?;
        let defined_name = match args.get(0)?.match_list() {
            Some(list) => list.get(0)?,
            _ => &args[0]
        };
        let tle_name = defined_name.match_atom()?;
        Some((tle_name.clone(), defined_name.id, defined_name))
    }
}

pub struct TopLevelExpressionIndex {
    expr_index: usize,
    atom_index: u64
}

struct Graph {
    adjacency_list: Vec<Vec<usize>>
}

impl Graph {
    fn new() -> Self {
        Self { adjacency_list: Vec::new() }
    }

    fn add_node(&mut self, _expr_index: usize) {
        self.adjacency_list.push(vec![]);
    }

    fn add_directed_edge(&mut self, src_expr_index: usize, dst_expr_index: usize) {
        let list = self.adjacency_list.get_mut(src_expr_index).unwrap();
        list.push(dst_expr_index);
    }
    
    fn get_node_descendants(&self, expr_index: usize) -> Vec<usize> {
        self.adjacency_list[expr_index].clone()
    }

    fn has_node_descendants(&self, expr_index: usize) -> bool {
        self.adjacency_list[expr_index].len() > 0
    }

    fn nodes_count(&self) -> usize {
        self.adjacency_list.len()
    }

    fn edges_count(&self) -> ParseResult<u64> {
        let mut total: u64 = 0;
        for node in self.adjacency_list.iter() {
            total = total.checked_add(node.len() as u64)
                .ok_or_else(|| ParseErrors::CostOverflow)?;
        }
        Ok(total)
    }
}

struct GraphWalker {
    seen: HashSet<usize>,
}

impl GraphWalker {

    fn new() -> Self { Self { seen: HashSet::new() } }

    /// Depth-first search producing a post-order sort
    fn get_sorted_dependencies(&mut self, graph: &Graph) -> ParseResult<Vec<usize>> {
        let mut sorted_indexes = Vec::<usize>::new();
        for expr_index in 0..graph.nodes_count() {
            self.sort_dependencies_recursion(expr_index, graph, &mut sorted_indexes);
        }

        Ok(sorted_indexes)
    }

    fn sort_dependencies_recursion(&mut self, tle_index: usize, graph: &Graph, branch: &mut Vec<usize>) {
        if self.seen.contains(&tle_index) {
            return
        }

        self.seen.insert(tle_index);
        if let Some(list) = graph.adjacency_list.get(tle_index) {
            for neighbor in list.iter() {
                self.sort_dependencies_recursion(neighbor.clone(), graph, branch);
            }
        }
        branch.push(tle_index);
    }

    fn get_cycling_dependencies(&mut self, graph: &Graph, sorted_indexes: &Vec<usize>) -> Option<Vec<usize>> {
        let mut tainted: HashSet<usize> = HashSet::new();

        for node in sorted_indexes.iter() {
            let mut tainted_descendants_count = 0;
            let descendants = graph.get_node_descendants(*node);
            for descendant in descendants.iter() {
                if !graph.has_node_descendants(*descendant) || tainted.contains(descendant) {
                    tainted.insert(*descendant);
                    tainted_descendants_count += 1;
                }
            }
            if tainted_descendants_count == descendants.len() {
                tainted.insert(*node);
            }
        }

        if tainted.len() == sorted_indexes.len() {
            return None
        }

        let nodes = HashSet::from_iter(sorted_indexes.iter().cloned());
        let deps = nodes.difference(&tainted).map(|i| *i).collect();        
        Some(deps) 
    }
}

