{-# LANGUAGE FlexibleInstances #-}
module Beekeeper.Core.DeclarationTraversal where

import Beekeeper.Core.Identifiers (modulePrefixes, Location(..), PythonTree(..))
import Beekeeper.Core.Declaration (Declaration(..), DeclarationTree)

import Data.Maybe (mapMaybe)

import qualified Language.Python.Common.AST as Py
import Language.Python.Common.SrcLocation (SrcSpan)

class HasDeclarations py where
  declarations :: py -> [DeclarationTree]

-- Unfortunately, language-python uses Py.Expr for targets, which means we have
-- to separate out assignment targets and normal expressions on our end. So,
-- whenever we are passed a [Py.Expr annot] instance that we'd like to treat as
-- an assignment, we have to wrap it in its newtyped target clothing
newtype TargetList = TargetList [Target]
newtype Target = Target (Py.ExprSpan)

toTargets :: [Py.ExprSpan] -> TargetList
toTargets = TargetList . map Target

collect :: (HasDeclarations py) => (expr -> [py]) -> expr -> [DeclarationTree]
collect proj = concatMap declarations . proj

project :: (HasDeclarations py) => [expr -> py] -> expr -> [DeclarationTree]
project projs expr = concat [declarations (proj expr) | proj <- projs]

singleIdent :: (String -> Declaration) -> Py.IdentSpan -> DeclarationTree
singleIdent decType ident = SourceTree {sourceNode = decType . Py.ident_string $ ident,
                                        children   = [],
                                        location   = Span . Py.annot $ ident}

identNode :: (String -> Declaration) -> Py.Ident SrcSpan -> [DeclarationTree]
identNode decType = return . singleIdent decType

instance (HasDeclarations py) => HasDeclarations (Maybe py) where
  declarations = maybe [] declarations

instance (HasDeclarations py, HasDeclarations py') => HasDeclarations (py, py')
  where declarations (k, v) = declarations k ++ declarations v

instance HasDeclarations Py.ArgumentSpan where
  declarations = declarations . Py.arg_expr

instance HasDeclarations Py.ExprSpan where
  declarations expr@(Py.Var {}) = [SourceTree {sourceNode = node,
                                               children   = [],
                                               location   = Span . Py.expr_annot $ expr}]
    where node = Dereference . Py.ident_string . Py.var_ident $ expr

  declarations expr@(Py.Call {}) = fun_dec ++ args_dec
    where fun_dec = declarations . Py.call_fun $ expr
          args_dec = collect Py.call_args expr

  declarations expr@(Py.Subscript {}) = project [Py.subscriptee, Py.subscript_expr] expr

  declarations expr@(Py.SlicedExpr {}) = slicee ++ slices
    where slicee = declarations. Py.slicee $ expr
          slices = collect Py.slices expr

  declarations expr@(Py.CondExpr {}) = project projs expr
    where projs = [Py.ce_true_branch, Py.ce_condition, Py.ce_false_branch]

  -- Attribute accesses are handled as binary ops with op type Dot, so we don't
  -- recurse down the right side for that particular case
  declarations expr@(Py.BinaryOp {}) = case Py.operator expr of
    Py.Dot {} -> declarations . Py.left_op_arg $ expr
    _ -> project [Py.left_op_arg, Py.right_op_arg] expr

  declarations expr@(Py.UnaryOp {}) = declarations . Py.op_arg $ expr

  declarations expr@(Py.Lambda {}) = [node]
    where args = collect Py.lambda_args expr
          body = declarations . Py.lambda_body $ expr
          node = SourceTree {sourceNode = LambdaDec,
                             children   = args ++ body,
                             location   = Span . Py.expr_annot $ expr}

  declarations expr@(Py.Tuple {}) = collect Py.tuple_exprs expr
  declarations expr@(Py.Yield {}) = declarations . Py.yield_expr $ expr
  declarations expr@(Py.Generator {}) = declarations . Py.gen_comprehension $ expr
  declarations expr@(Py.ListComp {}) = declarations . Py.list_comprehension $ expr
  declarations expr@(Py.List {}) = collect Py.list_exprs expr
  declarations expr@(Py.Dictionary {}) = collect Py.dict_mappings expr
  declarations expr@(Py.DictComp {}) = declarations . Py.dict_comprehension $ expr
  declarations expr@(Py.Set {}) = collect Py.set_exprs expr
  declarations expr@(Py.SetComp {}) = declarations . Py.set_comprehension $ expr
  declarations expr@(Py.Paren {}) = declarations . Py.paren_expr $ expr

  declarations _ = []

instance (HasDeclarations py) => HasDeclarations (Py.ComprehensionSpan py) where
  declarations comp = exprDec ++ forDec
    where exprDec = declarations . Py.comprehension_expr $ comp
          forDec = declarations . Py.comprehension_for $ comp

instance HasDeclarations Py.CompIfSpan where
  declarations comp = ifDec ++ next
    where ifDec = declarations . Py.comp_if $ comp
          next  = declarations . Py.comp_if_iter $ comp

instance HasDeclarations Py.CompForSpan where
  declarations expr = target_decs ++ in_decs ++ next
    where target_decs = declarations . toTargets . Py.comp_for_exprs $ expr
          in_decs = declarations . Py.comp_in_expr $ expr
          next = declarations . Py.comp_for_iter $ expr

instance HasDeclarations Py.CompIterSpan where
  declarations (Py.IterFor {Py.comp_iter_for = expr}) = declarations expr
  declarations (Py.IterIf {Py.comp_iter_if = expr}) = declarations expr

instance HasDeclarations Py.DecoratorSpan where
  declarations dec = nameNode : collect Py.decorator_args dec
    where name = Py.ident_string . head . Py.decorator_name $ dec
          nameNode = SourceTree {sourceNode = Dereference name,
                                 children   = [],
                                 location   = Span . Py.decorator_annot $ dec}

instance HasDeclarations Py.ExceptClauseSpan where
  declarations exc = exn_component fst ++ exn_component snd
    where exn_component proj = declarations . fmap proj . Py.except_clause $ exc

instance HasDeclarations Py.HandlerSpan where
  declarations handler = clause_decs ++ body_decs
    where clause_decs = declarations . Py.handler_clause $ handler
          body_decs = collect Py.handler_suite handler

instance HasDeclarations Py.ImportItemSpan where
  declarations py_import = map (tree . node) sourceNamePairs
    where sourceNamePairs = case Py.import_as_name py_import of
            Just ident -> [(itemNames, Py.ident_string ident)]
            Nothing -> modulePrefixes itemNames
          itemNames = map Py.ident_string . Py.import_item_name $ py_import
          node = uncurry ImportModule
          tree dec = SourceTree {sourceNode = dec,
                                 children   = [],
                                 location   = Span . Py.import_item_annot $ py_import}

instance HasDeclarations Py.ParameterSpan where
  declarations (Py.Param {Py.param_name = ident}) = identNode LocalVar ident
  declarations (Py.VarArgsPos {Py.param_name = ident}) = identNode LocalVar ident
  declarations (Py.VarArgsKeyword {Py.param_name = ident}) = identNode LocalVar ident
  declarations (Py.UnPackTuple {Py.param_unpack_tuple = pt}) = declarations pt
  declarations (Py.EndPositional {}) = []

instance HasDeclarations Py.ParamTupleSpan where
  declarations t@(Py.ParamTupleName {}) = identNode LocalVar . Py.param_tuple_name $ t
  declarations (Py.ParamTuple {Py.param_tuple = pt}) = concatMap declarations pt

instance HasDeclarations Py.RaiseExprSpan where
  declarations (Py.RaiseV3 exn_args) = declarations exn_args
  declarations (Py.RaiseV2 exn_args) = declarations exn_args

instance HasDeclarations Py.SliceSpan where
  declarations slice@(Py.SliceProper {}) = lower ++ upper ++ stride
    where lower = declarations . Py.slice_lower $ slice
          upper = declarations . Py.slice_upper $ slice
          stride = declarations . Py.slice_stride $ slice

  declarations slice@(Py.SliceExpr {}) = declarations . Py.slice_expr $ slice

  declarations (Py.SliceEllipsis {}) = []

instance HasDeclarations Py.StatementSpan where
  declarations (Py.Import {Py.import_items = items}) = concatMap declarations items

  -- TODO: handle dots in some way (we'll do absolute vs. relative in another pass)
  declarations stmt@(Py.FromImport {}) = if hasDots
                                         then []
                                         else map makeTree nodes
    where hasDots = 0 < (Py.import_relative_dots . Py.from_module $ stmt)
          sourceModule = case Py.import_relative_module . Py.from_module $ stmt of
            Just module_source -> map Py.ident_string module_source
            Nothing -> []

          nodes = case Py.from_items stmt of
            Py.ImportEverything {} -> [StarImport {source = sourceModule}]
            fi@(Py.FromItems {}) -> map fromItemNode . Py.from_items_items $ fi

          makeTree node = SourceTree {sourceNode = node,
                                      children = [],
                                      location = Span . Py.stmt_annot $ stmt}

          fromItemNode item = ImportVar (sourceModule ++ [sourceName]) name
            where sourceName = Py.ident_string . Py.from_item_name $ item
                  name = case Py.from_as_name item of
                    Just as_name -> Py.ident_string as_name
                    Nothing -> sourceName

  declarations stmt@(Py.While {}) = cond_decs ++ body_decs ++ else_decs
    where cond_decs = declarations . Py.while_cond $ stmt
          body_decs = collect Py.while_body stmt
          else_decs = collect Py.while_else stmt

  declarations stmt@(Py.For {}) = target_decs ++ gen_decs ++ body_decs ++ else_decs
    where target_decs = declarations . toTargets . Py.for_targets $ stmt
          gen_decs = declarations . Py.for_generator $ stmt
          body_decs = collect Py.for_body stmt
          else_decs = collect Py.for_else stmt

  declarations fun@(Py.Fun {}) = [root {children = args ++ funChildren}]
    where root = singleIdent FunctionDec . Py.fun_name $ fun
          args = collect Py.fun_args fun
          funChildren = concatMap declarations . Py.fun_body $ fun

  declarations cls@(Py.Class {}) = [root {children = args ++ clsChildren}]
    where root = singleIdent ClassDec . Py.class_name $ cls
          args = collect Py.class_args cls
          clsChildren = concatMap declarations . Py.class_body $ cls

  declarations stmt@(Py.Conditional {}) = cond_decs ++ body_decs ++ else_decs
    where cond_decs = Py.cond_guards stmt >>= declarations . fst
          body_decs = Py.cond_guards stmt >>= snd >>= declarations
          else_decs = collect Py.cond_else stmt

  declarations stmt@(Py.Assign {}) = expr ++ targets
    where targets = declarations . toTargets . Py.assign_to $ stmt
          expr = declarations . Py.assign_expr $ stmt

  declarations stmt@(Py.AugmentedAssign {}) = expr ++ target
    where target = declarations . toTargets . return . Py.aug_assign_to $ stmt
          expr = declarations . Py.aug_assign_expr $ stmt

  declarations stmt@(Py.Decorated {}) = decorators ++ fun
    where decorators = collect Py.decorated_decorators stmt
          fun = declarations . Py.decorated_def $ stmt

  declarations stmt@(Py.Return {}) = declarations . Py.return_expr $ stmt

  declarations stmt@(Py.Try {}) = all_decs
    where body_decs = collect Py.try_body stmt
          except_decs = collect Py.try_excepts stmt
          else_decs = collect Py.try_else stmt
          finally_decs = collect Py.try_finally stmt
          all_decs = concat [body_decs, except_decs, else_decs, finally_decs]

  declarations stmt@(Py.Raise {}) = declarations . Py.raise_expr $ stmt

  declarations stmt@(Py.With {}) = context_vars ++ context_decs ++ body_decs
    where body_decs = collect Py.with_body stmt
          context_vars = concatMap (declarations . fst) . Py.with_context $ stmt
          context_targets = mapMaybe snd . Py.with_context $ stmt
          context_decs = declarations . toTargets $ context_targets

  declarations stmt@(Py.Delete {}) = collect Py.del_exprs stmt

  declarations stmt@(Py.StmtExpr {}) = declarations . Py.stmt_expr $ stmt

  declarations stmt@(Py.Global {Py.global_vars = vars}) = map makeTree vars
    where makeTree var = SourceTree {sourceNode = GlobalVar . Py.ident_string $ var,
                                     children   = [],
                                     location   = Span . Py.stmt_annot $ stmt}

  declarations stmt@(Py.NonLocal {Py.nonLocal_vars = vars}) = map makeTree vars
    where makeTree var = SourceTree {sourceNode = NonlocalVar . Py.ident_string $ var,
                                     children   = [],
                                     location   = Span . Py.stmt_annot $ stmt}

  declarations stmt@(Py.Assert {}) = collect Py.assert_exprs stmt

  declarations stmt@(Py.Print {}) = collect Py.print_exprs stmt

  declarations stmt@(Py.Exec {}) = execExpr ++ scopes
    where execExpr = declarations . Py.exec_expr $ stmt
          scopes = declarations . Py.exec_globals_locals $ stmt

  declarations _ = []

-- TODO: This handles most cases but not the complete grammar
instance HasDeclarations Target where
  declarations (Target expr) = case expr of
    Py.Var {Py.var_ident = ident} -> identNode LocalVar ident
    Py.Paren {Py.paren_expr = p} -> declarations . toTargets . Py.tuple_exprs $ p
    Py.Tuple {Py.tuple_exprs = t} -> declarations . toTargets $ t
    Py.List {Py.list_exprs = l} -> declarations . toTargets $ l
    Py.BinaryOp {Py.operator = Py.Dot {}} -> declarations . Py.left_op_arg $ expr
    Py.Subscript {} -> declarations expr
    Py.SlicedExpr {} -> declarations expr
    Py.Starred {} -> []
    _ -> error "broken assumption: bad Target Py.Expr"

instance HasDeclarations TargetList where
  declarations (TargetList targets) = concatMap declarations targets