use std::collections::{HashMap, HashSet};

use vm::analysis::AnalysisDatabase;
use vm::representations::{SymbolicExpression, PreSymbolicExpression, ClarityName, TraitDefinition};
use vm::types::{Value, TraitIdentifier, QualifiedContractIdentifier};
use vm::functions::NativeFunctions;
use vm::functions::define::{DefineFunctions, DefineFunctionsParsed};
use vm::ast::types::{ContractAST, BuildASTPass, PreExpressionsDrain};
use vm::ast::errors::{ParseResult, ParseError, ParseErrors};
use vm::representations::PreSymbolicExpressionType::{AtomValue, Atom, List, TraitReference, SugaredFieldIdentifier, FieldIdentifier};            

pub struct TraitsResolver {
}

impl BuildASTPass for TraitsResolver {

    fn run_pass(contract_ast: &mut ContractAST) -> ParseResult<()> {
        let mut command = TraitsResolver::new();
        command.run(contract_ast)?;
        Ok(())
    }
}

impl TraitsResolver {

    fn new() -> TraitsResolver {
        TraitsResolver {}
    }

    pub fn run(&mut self, contract_ast: &mut ContractAST) -> ParseResult<()> {
        let exprs = contract_ast.pre_expressions[..].to_vec();
        let mut referenced_traits = HashMap::new();

        for exp in exprs.iter() {
            
            let (define_type, args) = match self.try_parse_pre_expr(exp) {
                Some(x) => x,
                None => continue 
            };

            match define_type {
                DefineFunctions::Trait => {
                    if args.len() != 2 {
                        return Err(ParseErrors::DefineTraitBadSignature.into())
                    }

                    match (&args[0].pre_expr, &args[1].pre_expr) {
                        (Atom(trait_name), List(trait_definition)) => {
                            // Check for collisions
                            if contract_ast.referenced_traits.contains_key(trait_name) {
                                return Err(ParseErrors::NameAlreadyUsed(trait_name.to_string()).into())
                            }

                            // Traverse and probe for generics nested in the trait definition
                            self.probe_for_generics(trait_definition, &mut referenced_traits, true)?;

                            let trait_id = TraitIdentifier {
                                name: trait_name.clone(),
                                contract_identifier: contract_ast.contract_identifier.clone()      
                            };
                            contract_ast.referenced_traits.insert(trait_name.clone(), TraitDefinition::Defined(trait_id));
                        },
                        _ => return Err(ParseErrors::DefineTraitBadSignature.into())
                    }
                },
                DefineFunctions::UseTrait => {
                    if args.len() != 2 {
                        return Err(ParseErrors::ImportTraitBadSignature.into())
                    }

                    if let Some(trait_name) = args[0].match_atom() {
                        // Check for collisions
                        if contract_ast.referenced_traits.contains_key(trait_name) {
                            return Err(ParseErrors::NameAlreadyUsed(trait_name.to_string()).into())
                        }

                        let trait_id = match &args[1].pre_expr {
                            SugaredFieldIdentifier(contract_name, name) => {
                                let contract_identifier = QualifiedContractIdentifier::new(
                                    contract_ast.contract_identifier.issuer.clone(), 
                                    contract_name.clone());
                                TraitIdentifier { name: name.clone(), contract_identifier}
                            },
                            FieldIdentifier(trait_identifier) => trait_identifier.clone(),
                            _ => return Err(ParseErrors::ImportTraitBadSignature.into()),
                        };
                        contract_ast.referenced_traits.insert(trait_name.clone(), TraitDefinition::Imported(trait_id));
                    } else {
                        return Err(ParseErrors::ImportTraitBadSignature.into());
                    }
                },
                DefineFunctions::ImplTrait => {
                    if args.len() != 1 {
                        return Err(ParseErrors::ImplTraitBadSignature.into())
                    }

                    let trait_id = match &args[0].pre_expr {
                        SugaredFieldIdentifier(contract_name, name) => {
                            let contract_identifier = QualifiedContractIdentifier::new(
                                contract_ast.contract_identifier.issuer.clone(), 
                                contract_name.clone());
                            TraitIdentifier { name: name.clone(), contract_identifier}
                        },
                        FieldIdentifier(trait_identifier) => trait_identifier.clone(),
                        _ => return Err(ParseErrors::ImplTraitBadSignature.into()),
                    };
                    contract_ast.implemented_traits.insert(trait_id);    
                },
                DefineFunctions::PublicFunction | DefineFunctions::PrivateFunction | DefineFunctions::ReadOnlyFunction => {
                    // Traverse and probe for generics in functions type definitions
                    self.probe_for_generics(&args, &mut referenced_traits, true)?;
                },
                DefineFunctions::Constant | DefineFunctions::Map | DefineFunctions::PersistedVariable | 
                DefineFunctions::FungibleToken | DefineFunctions::NonFungibleToken => {
                    self.probe_for_generics(&args[1..], &mut referenced_traits, false)?;
                }
            };
        }

        for (trait_reference, expr) in referenced_traits {
            if !contract_ast.referenced_traits.contains_key(&trait_reference) {
                let mut err = ParseError::new(ParseErrors::TraitReferenceUnknown(trait_reference.to_string()));
                err.set_pre_expression(&expr);
                return Err(err.into());
            }
        }

        Ok(())
    }

    fn try_parse_pre_expr<'a>(&self, expression: &'a PreSymbolicExpression) -> Option<(DefineFunctions, &'a [PreSymbolicExpression])> {
        let expression = expression.match_list()?;
        let (function_name, args) = expression.split_first()?;
        let function_name = function_name.match_atom()?;
        let define_type = DefineFunctions::lookup_by_name(function_name)?;
        Some((define_type, args))
    }

    fn probe_for_generics(&mut self, exprs: &[PreSymbolicExpression], 
                          referenced_traits: &mut HashMap<ClarityName, PreSymbolicExpression>, 
                          should_reference: bool) -> ParseResult<()>  {
        for expression in exprs.iter() {

            match &expression.pre_expr {
                List(list) => { self.probe_for_generics(&list, referenced_traits, should_reference)?; },
                TraitReference(trait_name) => {
                    if should_reference {
                        referenced_traits.insert(trait_name.clone(), expression.clone());
                    } else {
                        return Err(ParseErrors::TraitReferenceNotAllowed.into())
                    }
                },
                _ => { /* no-op */ }
            }
        }
        Ok(())
    }
}
