# This class supports type checking. Beancount code dynamically sets the

# ``store`` attribute, in bc_query_execute.execute_query().

class Context(bc_query_execute.RowContext):

    store: Store

class MetaDocs(bc_query_env.AnyMeta):

    """Return a list of document links from metadata."""

    def __init__(self, operands: List[str]) -> None:

        super(bc_query_env.AnyMeta, self).__init__(operands, list)

        # The second argument is our return type.

        # It should match the annotated return type of __call__.

    def __call__(self, context: Context) -> List[str]:

        raw_value = super().__call__(context)

        if isinstance(raw_value, str):

            return raw_value.split()

        else:

            return []

class StrMeta(bc_query_env.AnyMeta):

    """Looks up metadata like AnyMeta, then always returns a string."""

    def __init__(self, operands: List[str]) -> None:

        super(bc_query_env.AnyMeta, self).__init__(operands, str)

    def __call__(self, context: Context) -> str:

        raw_value = super().__call__(context)

        if raw_value is None:

            return ''

        else:

            return str(raw_value)

class AggregateSet(bc_query_compile.EvalAggregator):

    __intypes__ = [object]

    def __init__(self, operands: List[str]) -> None:

       super().__init__(operands, set)

    def allocate(self, allocator: bc_query_execute.Allocator) -> None:

        """Allocate and save an index handle into result storage."""

        self.handle = allocator.allocate()

    def initialize(self, store: Store) -> None:

        """Prepare result storage for a new aggregation."""

        store[self.handle] = self.dtype()

        # self.dtype() is our return type, aka the second argument to __init__

        # above, aka the annotated return type of __call__.

    def update(self, store: Store, context: Context) -> None:

        """Update existing storage with new result data."""

        value, = self.eval_args(context)

        if isinstance(value, Sequence) and not isinstance(value, (str, tuple)):

            store[self.handle].update(value)

        else:

            store[self.handle].add(value)

    def __call__(self, context: Context) -> set:

        """Return the result for an aggregation."""

        return context.store[self.handle]  # type:ignore[no-any-return]

class FilterPostingsEnvironment(bc_query_env.FilterPostingsEnvironment):

    functions: EnvironmentFunctions = bc_query_env.FilterPostingsEnvironment.functions.copy()  # type:ignore[assignment]

    functions['meta_docs'] = MetaDocs

    functions['str_meta'] = StrMeta

class TargetsEnvironment(bc_query_env.TargetsEnvironment):

    functions = FilterPostingsEnvironment.functions.copy()

    functions.update(bc_query_env.AGGREGATOR_FUNCTIONS)

    functions['set'] = AggregateSet

class BQLShell(bc_query_shell.BQLShell):

    def __init__(

            self,

            is_interactive: bool,

            loadfun: Callable[[], books.LoadResult],

            outfile: TextIO,

            default_format: str='text',

            do_numberify: bool=False,

            rt_wrapper: Optional[rtutil.RT]=None,

    ) -> None:

        super().__init__(is_interactive, loadfun, outfile, default_format, do_numberify)

        self.env_postings = FilterPostingsEnvironment()

        self.env_targets = TargetsEnvironment()

        self.ods = QueryODS(rt_wrapper)

    def on_Select(self, statement: QueryStatement) -> None:

        output_format: str = self.vars['format']

        try:

            render_func = getattr(self, f'_render_{output_format}')

        except AttributeError:

            logger.error("unknown output format %r", output_format)

            return

        try:

            logger.debug("compiling query")

            compiled_query = bc_query_compile.compile(

                statement, self.env_targets, self.env_postings, self.env_entries,

            )

            logger.debug("executing query")

            row_types, rows = bc_query_execute.execute_query(

                compiled_query, self.entries, self.options_map,

            )

            if self.vars['numberify']:

                logger.debug("numberifying query")

                row_types, rows = bc_query_numberify.numberify_results(

                    row_types, rows, self.options_map['dcontext'].build(),

                )

        except Exception as error:

            logger.error(str(error), exc_info=logger.isEnabledFor(logging.DEBUG))

            return

        if not rows and output_format != 'ods':

            print("(empty)", file=self.outfile)

        else:

            logger.debug("rendering query as %s", output_format)

            render_func(statement, row_types, rows)

    def _render_csv(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None:

        bc_query_render.render_csv(

            row_types,

            rows,

            self.options_map['dcontext'],

            self.outfile,

            self.vars['expand'],

        )

    def _render_ods(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None:

        self.ods.write_query(statement, row_types, rows)

        logger.info(

            "%s rows of results saved in sheet %s",

            len(rows),

            self.ods.sheet.getAttribute('name'),

        )

    def _render_text(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None:

        with contextlib.ExitStack() as stack:

            if self.is_interactive:

                output = stack.enter_context(self.get_pager())

            else:

                output = self.outfile

            bc_query_render.render_text(

                row_types,

                rows,

                self.options_map['dcontext'],

                output,

                self.vars['expand'],

                self.vars['boxed'],

                self.vars['spaced'],

            )

        # Names of functions, as defined in Environments, that look up

        # posting metadata that could contain documentation links

        # The configuration hasn't been loaded, so we don't know the boundaries

        # of a fiscal year yet. But that's okay, because we just need to set

        # enough so that when these arguments are passed to a BooksLoader,

        # it'll load the right fiscal year.