module Beekeeper.Pretty where

import Beekeeper.Core.Declaration (Declaration(..))
import Beekeeper.Core.Identifiers (ImportSource, ModuleName(..), PythonTree(..))

import qualified Data.List as List

import Text.PrettyPrint (
  Doc,
  colon,
  empty,
  int,
  nest,
  parens,
  text,
  vcat,
  (<>),
  (<+>),
  ($+$))

class Pretty a where
  pretty :: a -> Doc

instance (Pretty a, Pretty b) => Pretty (PythonTree a b) where
  pretty tree = header <+> countInfo $+$ prettyChildren
    where header = case tree of
            ModuleTree {} -> pretty . moduleNode $ tree
            SourceTree {} -> pretty . sourceNode $ tree
          count = length . children $ tree
          annot | count == 1 = "child"
                | otherwise  = "children"
          countInfo | count > 0 = parens (int count <+> text annot) <> colon
                    | otherwise = empty
          prettyChildren = vcat . map (nest 4 . pretty) . children $ tree

instance Pretty ModuleName where
  pretty m@(ModuleName {}) = text . moduleName $ m
  pretty m@(SourceName {}) = text $ moduleName m ++ ".py"

instance Pretty Declaration where
  pretty dec@(StarImport {}) = text "from" <+> src <+> text "import *"
    where src = text . sourceToName . source $ dec

  pretty dec@(ImportVar {}) = text "import" <+> src <+> text "as" <+> name
    where src = text . sourceToName . source $ dec
          name = text . declarationName $ dec

  pretty dec@(ImportModule {}) = text "import" <+> src <+> text "as" <+> name
    where src = text . sourceToName . source $ dec
          name = text . declarationName $ dec

  pretty (LambdaDec {}) = text "lambda"

  pretty dec = header <+> varName
    where header = text $ case dec of
            ClassDec {}    -> "class"
            FunctionDec {} -> "def"
            LocalVar {}    -> "var"
            GlobalVar {}   -> "global"
            NonlocalVar {} -> "nonlocal"
            Dereference {} -> "load"
            _              -> error "bad declaration pretty print instance"

          varName = text . declarationName $ dec

sourceToName :: ImportSource -> String
sourceToName = List.intercalate "."