"""query.py - Report arbitrary queries with advanced loading and formatting

This tool extends Beancount's bean-query with the following:

* ``META_DOCS("meta-key")`` function: Given the name of metadata with
  documentation in it, returns a list of the documentation links.

* ``STR_META("meta-key")`` function: Looks up metadata like ANY_META, then
  returns the result as a nicely formatted string, to make sorting easier or
  just provide prettier reports.

* ``SET()`` function: Returns only unique arguments over a GROUP BY.

* ODS output format, with proper formatting of currency and documentation links.

* Loads books from your configuration file, with options available to specify
  which year(s) to load.

* Can load rewrite rules just like other reports. Your queries will show the
  transformed entries.

* Improved error reporting.

Run it like bean-query, except instead of specifying a books filename, use
date options like ``--begin``, ``--end``, ``--fy``, and ``--cy`` to specify
the year(s) you want to load. Run ``query-report --help`` for details about
those.

Start an interactive shell::

    query-report [year options]

Write a spreadsheet with results for one query::

    query-report [year options] [-O OUTPUT.ods] <query string>

    query-report [year options] [-O OUTPUT.ods] < QUERY_FILE.bql

query-report also accepts all the same options as bean-query, like ``--format``
and ``--numberify``.
"""
# Copyright © 2021  Brett Smith
# License: AGPLv3-or-later WITH Beancount-Plugin-Additional-Permission-1.0
#
# Full copyright and licensing details can be found at toplevel file
# LICENSE.txt in the repository.

import argparse
import contextlib
import datetime
import enum
import functools
import itertools
import logging
import sys

from typing import (
    cast,
    Any,
    Callable,
    Dict,
    Iterable,
    Iterator,
    List,
    Mapping,
    NamedTuple,
    Optional,
    Sequence,
    TextIO,
    Tuple,
    Type,
    Union,
)
from ..beancount_types import (
    MetaKey,
    MetaValue,
)

from decimal import Decimal
from pathlib import Path
from beancount.core.amount import _Amount as BeancountAmount
from beancount.core.inventory import Inventory
from beancount.core.position import _Position as Position

import beancount.query.numberify as bc_query_numberify
import beancount.query.query_compile as bc_query_compile
import beancount.query.query_env as bc_query_env
import beancount.query.query_execute as bc_query_execute
import beancount.query.query_parser as bc_query_parser
import beancount.query.query_render as bc_query_render
import beancount.query.shell as bc_query_shell
import odf.table  # type:ignore[import]

from . import core
from . import rewrite
from .. import books
from .. import cliutil
from .. import config as configmod
from .. import data
from .. import rtutil

PROGNAME = 'query-report'
logger = logging.getLogger('conservancy_beancount.reports.query')

CellFunc = Callable[[Any], odf.table.TableCell]
EnvironmentFunctions = Dict[
    # The real key type is something like:
    #   Union[str, Tuple[str, Type, ...]]
    # but two issues with that. One, you can't use Ellipses in a Tuple like
    # that, so there's no short way to declare this. Second, Beancount doesn't
    # declare it anyway, and mypy infers it as Sequence[object]. So just use
    # that.
    Sequence[object],
    Type[bc_query_compile.EvalFunction],
]
RowTypes = Sequence[Tuple[str, Type]]
Rows = Sequence[NamedTuple]
Store = List[Any]
QueryExpression = Union[
    bc_query_parser.Column,
    bc_query_parser.Constant,
    bc_query_parser.Function,
    bc_query_parser.UnaryOp,
]
QueryStatement = Union[
    bc_query_parser.Balances,
    bc_query_parser.Journal,
    bc_query_parser.Select,
]

# 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):
    """Filter argument values that aren't unique."""
    __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 BooksLoader:
    """Closure to load books with a zero-argument callable

    This matches the load interface that BQLShell expects.
    """
    def __init__(
            self,
            books_loader: Optional[books.Loader],
            start_date: Optional[datetime.date]=None,
            stop_date: Optional[datetime.date]=None,
            rewrite_rules: Sequence[rewrite.RewriteRuleset]=(),
    ) -> None:
        self.books_loader = books_loader
        self.start_date = start_date
        self.stop_date = stop_date
        self.rewrite_rules = rewrite_rules

    def __call__(self) -> books.LoadResult:
        logger.debug("BooksLoader called")
        result = books.Loader.dispatch(self.books_loader, self.start_date, self.stop_date)
        logger.debug("books loaded from Beancount")
        if self.rewrite_rules:
            for index, entry in enumerate(result.entries):
                # entry might not be a Transaction; we catch that later.
                # The type ignores are because the underlying Beancount type isn't
                # type-checkable.
                postings = data.Posting.from_txn(entry)  # type:ignore[arg-type]
                for ruleset in self.rewrite_rules:
                    postings = ruleset.rewrite(postings)
                try:
                    result.entries[index] = entry._replace(postings=list(postings))  # type:ignore[call-arg]
                except AttributeError:
                    pass
            logger.debug("rewrite rules applied")
        return result


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)
        self.last_line_parsed = ''

    def run_parser(
            self,
            line: str,
            default_close_date: Optional[datetime.datetime]=None,
    ) -> None:
        self.last_line_parsed = line
        super().run_parser(line, default_close_date)

    @functools.wraps(bc_query_shell.BQLShell.on_Select, ('__doc__',))
    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))
            if (isinstance(error, TypeError)
                and error.args
                and ' not supported between instances ' in error.args[0]):
                logger.info(
                    "HINT: Are you using ORDER BY or comparisons with metadata "
                    "that isn't consistently set?\n  "
                    "Try looking up that metadata with str_meta() instead to "
                    "ensure your comparisons use a consistent data type.",
                )
            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, self.last_line_parsed)
        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'],
            )


class QueryODS(core.BaseODS[NamedTuple, None]):
    META_FNAMES = frozenset([
        # Names of functions, as defined in Environments, that look up
        # posting metadata that could contain documentation links
        'any_meta',
        'entry_meta',
        'meta',
        'meta_docs',
        'str_meta',
    ])

    def is_empty(self) -> bool:
        return not self.sheet.childNodes

    def section_key(self, row: NamedTuple) -> None:
        return None

    def _generic_cell(self, value: Any) -> odf.table.TableCell:
        if isinstance(value, Iterable) and not isinstance(value, (str, tuple)):
            return self.multiline_cell(value)
        else:
            return self.string_cell('' if value is None else str(value))

    def _inventory_cell(self, value: Inventory) -> odf.table.TableCell:
        return self.balance_cell(core.Balance(pos.units for pos in value))

    def _link_string_cell(self, value: str) -> odf.table.TableCell:
        return self.meta_links_cell(value.split())

    def _metadata_cell(self, value: MetaValue) -> odf.table.TableCell:
        return self._cell_type(type(value))(value)

    def _position_cell(self, value: Position) -> odf.table.TableCell:
        return self.currency_cell(value.units)

    def _cell_type(self, row_type: Type) -> CellFunc:
        """Return a function to create a cell, for non-metadata row types."""
        if issubclass(row_type, Inventory):
            return self._inventory_cell
        elif issubclass(row_type, Position):
            return self._position_cell
        elif issubclass(row_type, BeancountAmount):
            return self.currency_cell
        elif issubclass(row_type, (int, float, Decimal)):
            return self.float_cell
        elif issubclass(row_type, datetime.date):
            return self.date_cell
        elif issubclass(row_type, str):
            return self.string_cell
        else:
            return self._generic_cell

    def _link_cell_type(self, row_type: Type) -> CellFunc:
        """Return a function to create a cell from metadata with documentation links."""
        if issubclass(row_type, str):
            return self._link_string_cell
        elif issubclass(row_type, tuple):
            return self._generic_cell
        elif issubclass(row_type, Iterable):
            return self.meta_links_cell
        else:
            return self._generic_cell

    def _meta_target(self, target: QueryExpression) -> Optional[MetaKey]:
        """Return the metadata key looked up by this target, if any

        This function takes a parsed target (i.e., what we're SELECTing) and
        recurses it to see whether it's looking up any metadata. If so, it
        returns the key of that metadata. Otherwise it returns None.
        """
        if isinstance(target, bc_query_parser.UnaryOp):
            return self._meta_target(target.operand)
        elif not isinstance(target, bc_query_parser.Function):
            return None
        try:
            operand = target.operands[0]
        except IndexError:
            return None
        if (target.fname in self.META_FNAMES
            and isinstance(operand, bc_query_parser.Constant)):
            return operand.value  # type:ignore[no-any-return]
        else:
            for operand in target.operands:
                retval = self._meta_target(operand)
                if retval is not None:
                    break
            return retval

    def _cell_types(self, statement: QueryStatement, row_types: RowTypes) -> Iterator[CellFunc]:
        """Return functions to create table cells from result rows

        Given a parsed query and the types of return rows, yields a function
        to create a cell for each column in the row, in order. The returned
        functions vary in order to provide the best available formatting for
        different data types.
        """
        if (isinstance(statement, bc_query_parser.Select)
            and isinstance(statement.targets, Sequence)):
            targets = [t.expression for t in statement.targets]
        else:
            # Synthesize something that makes clear we're not loading metadata.
            targets = [bc_query_parser.Column(name) for name, _ in row_types]
        for target, (_, row_type) in zip(targets, row_types):
            meta_key = self._meta_target(target)
            if meta_key is None:
                yield self._cell_type(row_type)
            elif meta_key in data.LINK_METADATA:
                yield self._link_cell_type(row_type)
            else:
                yield self._metadata_cell

    def write_query(
            self,
            statement: QueryStatement,
            row_types: RowTypes,
            rows: Rows,
            query_string: Optional[str]=None,
    ) -> None:
        if self.is_empty():
            self.sheet.setAttribute('name', "Query 1")
        else:
            self.use_sheet(f"Query {len(self.document.spreadsheet.childNodes) + 1}")
        for name, row_type in row_types:
            if issubclass(row_type, datetime.date):
                col_width = 1.0
            elif issubclass(row_type, (BeancountAmount, Inventory, Position)):
                col_width = 1.5
            else:
                col_width = 2.0
            col_style = self.column_style(col_width)
            self.sheet.addElement(odf.table.TableColumn(stylename=col_style))
        self.add_row(*(
            self.string_cell(data.Metadata.human_name(name), stylename=self.style_bold)
            for name, _ in row_types
        ))
        self.lock_first_row()
        if query_string:
            self.add_annotation(query_string, parent=self.sheet.lastChild.firstChild)
        cell_funcs = list(self._cell_types(statement, row_types))
        for row in rows:
            self.add_row(*(
                cell_func(value)
                for cell_func, value in zip(cell_funcs, row)
            ))


class ReportFormat(enum.Enum):
    TEXT = 'text'
    TXT = TEXT
    CSV = 'csv'
    ODS = 'ods'


class SetCYDates(argparse.Action):
    def __call__(self,
                 parser: argparse.ArgumentParser,
                 namespace: argparse.Namespace,
                 values: Union[Sequence[Any], str, None]=None,
                 option_string: Optional[str]=None,
    ) -> None:
        value = cliutil.year_or_date_arg(str(values))
        if isinstance(value, int):
            value = datetime.date(value, 1, 1)
        namespace.start_date = value
        namespace.stop_date = cliutil.diff_year(value, 1)


class SetFYDates(argparse.Action):
    def __call__(self,
                 parser: argparse.ArgumentParser,
                 namespace: argparse.Namespace,
                 values: Union[Sequence[Any], str, None]=None,
                 option_string: Optional[str]=None,
    ) -> None:
        value = cliutil.year_or_date_arg(str(values))
        namespace.start_date = value
        # 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.
        if isinstance(value, int):
            namespace.stop_date = value
        else:
            namespace.stop_date = value + datetime.timedelta(days=1)


def parse_arguments(arglist: Optional[Sequence[str]]=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(prog=PROGNAME)
    cliutil.add_version_argument(parser)
    cliutil.add_loglevel_argument(parser)
    parser.add_argument(
        '--begin', '--start', '-b',
        dest='start_date',
        metavar='YEAR',
        type=cliutil.year_or_date_arg,
        help="""Begin loading entries from this fiscal year. You can specify a
full date, and %(prog)s will use the fiscal year for that date.
""")
    parser.add_argument(
        '--end', '--stop', '-e',
        dest='stop_date',
        metavar='YEAR',
        type=cliutil.year_or_date_arg,
        help="""End loading entries at this fiscal year. You can specify a
full date, and %(prog)s will use the fiscal year for that date.
""")
    parser.add_argument(
        '--calendar-year', '--cy',
        action=SetCYDates,
        metavar='YEAR',
        help="""Shortcut to set --begin and --end to load a single calendar year.
You can specify a full date, or just a year to start from January 1.
""")
    parser.add_argument(
        '--fiscal-year', '--fy',
        action=SetFYDates,
        metavar='YEAR',
        help="""Shortcut to set --begin and --end to load a single fiscal year.
You can specify a full date, and %(prog)s will use the fiscal year for that date.
""")
    cliutil.add_rewrite_rules_argument(parser)
    format_arg = cliutil.EnumArgument(ReportFormat)
    parser.add_argument(
        '--report-type', '--format', '-t', '-f',
        metavar='TYPE',
        type=format_arg.enum_type,
        help="""Format of report to generate. Choices are
{format_arg.choices_str()}. Default is guessed from your output filename
extension. If that fails, default is 'text' for interactive output, and 'ods'
otherwise.
""")
    parser.add_argument(
        '--numberify', '-m',
        action='store_true',
        help="""Separate currency amounts into numeric columns by currency
""")
    parser.add_argument(
        '--output-file', '-O', '-o',
        metavar='PATH',
        type=Path,
        help="""Write the report to this file, or stdout when PATH is `-`.
The default is stdout for text and CSV reports, and a generated filename for
ODS reports.
""")
    parser.add_argument(
        'query',
        nargs=argparse.ZERO_OR_MORE,
        default=[],
        help="""Query to run non-interactively. If none is provided, and
standard input is not a terminal, reads the query from stdin instead.
""")
    return parser.parse_args(arglist)

def main(arglist: Optional[Sequence[str]]=None,
         stdout: TextIO=sys.stdout,
         stderr: TextIO=sys.stderr,
         config: Optional[configmod.Config]=None,
) -> int:
    args = parse_arguments(arglist)
    cliutil.set_loglevel(logger, args.loglevel)
    if config is None:
        config = configmod.Config()
        config.load_file()

    query = ' '.join(args.query).strip()
    if not query and not sys.stdin.isatty():
        query = sys.stdin.read().strip()
    if args.report_type is None:
        try:
            args.report_type = ReportFormat[args.output_file.suffix[1:].upper()]
        except (AttributeError, KeyError):
            args.report_type = ReportFormat.ODS if query else ReportFormat.TEXT

    load_func = BooksLoader(
        config.books_loader(),
        args.start_date,
        args.stop_date,
        [rewrite.RewriteRuleset.from_yaml(path) for path in args.rewrite_rules],
    )
    shell = BQLShell(
        not query,
        load_func,
        stdout,
        args.report_type.value,
        args.numberify,
        config.rt_wrapper(),
    )
    shell.on_Reload()
    if query:
        shell.onecmd(query)
    else:
        shell.cmdloop()

    if not shell.ods.is_empty():
        shell.ods.set_common_properties(config.books_repo())
        if args.output_file is None:
            out_dir_path = config.repository_path() or Path()
            args.output_file = out_dir_path / 'QueryResults_{}.ods'.format(
                datetime.datetime.now().isoformat(timespec='seconds'),
            )
            logger.info("Writing spreadsheet to %s", args.output_file)
        ods_file = cliutil.bytes_output(args.output_file, stdout)
        shell.ods.save_file(ods_file)

    return cliutil.ExitCode.OK

entry_point = cliutil.make_entry_point(__name__, PROGNAME)

if __name__ == '__main__':
    exit(entry_point())