"""Example importer for example broker UTrade."""

__copyright__ = "Copyright (C) 2016  Martin Blais"
__license__ = "GNU GPLv2"

import csv
import datetime
import re
import logging
from os import path

from dateutil.parser import parse

from beancount.core.number import D
from beancount.core.number import ZERO
from beancount.core import data
from beancount.core import account
from beancount.core import amount
from beancount.core import position
from beancount.ingest import importer


class Importer(importer.ImporterProtocol):
    """An importer for UTrade CSV files (an example investment bank)."""

    def __init__(
        self,
        currency,
        account_root,
        account_cash,
        account_dividends,
        account_gains,
        account_fees,
        account_external,
    ):
        self.currency = currency
        self.account_root = account_root
        self.account_cash = account_cash
        self.account_dividends = account_dividends
        self.account_gains = account_gains
        self.account_fees = account_fees
        self.account_external = account_external

    def identify(self, file):
        # Match if the filename is as downloaded and the header has the unique
        # fields combination we're looking for.
        return re.match(
            r"UTrade\d\d\d\d\d\d\d\d\.csv", path.basename(file.name)
        ) and re.match("DATE,TYPE,REF", file.head())

    def file_name(self, file):
        return "utrade.{}".format(path.basename(file.name))

    def file_account(self, _):
        return self.account_root

    def file_date(self, file):
        # Extract the statement date from the filename.
        return datetime.datetime.strptime(
            path.basename(file.name), "UTrade%Y%m%d.csv"
        ).date()

    def extract(self, file):
        # Open the CSV file and create directives.
        entries = []
        index = 0
        with open(file.name) as infile:
            for index, row in enumerate(csv.DictReader(infile)):
                meta = data.new_metadata(file.name, index)
                date = parse(row["DATE"]).date()
                rtype = row["TYPE"]
                link = "ut{0[REF #]}".format(row)
                desc = "({0[TYPE]}) {0[DESCRIPTION]}".format(row)
                units = amount.Amount(D(row["AMOUNT"]), self.currency)
                fees = amount.Amount(D(row["FEES"]), self.currency)
                other = amount.add(units, fees)

                if rtype == "XFER":
                    assert fees.number == ZERO
                    txn = data.Transaction(
                        meta,
                        date,
                        self.FLAG,
                        None,
                        desc,
                        data.EMPTY_SET,
                        {link},
                        [
                            data.Posting(self.account_cash, units, None, None, None, None),
                            data.Posting(
                                self.account_external, -other, None, None, None, None
                            ),
                        ],
                    )

                elif rtype == "DIV":
                    assert fees.number == ZERO

                    # Extract the instrument name from its description.
                    match = re.search(r"~([A-Z]+)$", row["DESCRIPTION"])
                    if not match:
                        logging.error("Missing instrument name in '%s'", row["DESCRIPTION"])
                        continue
                    instrument = match.group(1)
                    account_dividends = self.account_dividends.format(instrument)

                    txn = data.Transaction(
                        meta,
                        date,
                        self.FLAG,
                        None,
                        desc,
                        data.EMPTY_SET,
                        {link},
                        [
                            data.Posting(self.account_cash, units, None, None, None, None),
                            data.Posting(account_dividends, -other, None, None, None, None),
                        ],
                    )

                elif rtype in ("BUY", "SELL"):
                    # Extract the instrument name, number of units, and price from
                    # the description. That's just what we're provided with (this is
                    # actually realistic of some data from some institutions, you
                    # have to figure out a way in your parser).
                    match = re.search(
                        r"\+([A-Z]+)\b +([0-9.]+)\b +@([0-9.]+)", row["DESCRIPTION"]
                    )
                    if not match:
                        logging.error("Missing purchase infos in '%s'", row["DESCRIPTION"])
                        continue
                    instrument = match.group(1)
                    account_inst = account.join(self.account_root, instrument)
                    units_inst = amount.Amount(D(match.group(2)), instrument)
                    rate = D(match.group(3))

                    if rtype == "BUY":
                        cost = position.Cost(rate, self.currency, None, None)
                        txn = data.Transaction(
                            meta,
                            date,
                            self.FLAG,
                            None,
                            desc,
                            data.EMPTY_SET,
                            {link},
                            [
                                data.Posting(
                                    self.account_cash, units, None, None, None, None
                                ),
                                data.Posting(
                                    self.account_fees, fees, None, None, None, None
                                ),
                                data.Posting(
                                    account_inst, units_inst, cost, None, None, None
                                ),
                            ],
                        )

                    elif rtype == "SELL":
                        # Extract the lot. In practice this information not be there
                        # and you will have to identify the lots manually by editing
                        # the resulting output. You can leave the cost.number slot
                        # set to None if you like.
                        match = re.search(r"\(LOT ([0-9.]+)\)", row["DESCRIPTION"])
                        if not match:
                            logging.error("Missing cost basis in '%s'", row["DESCRIPTION"])
                            continue
                        cost_number = D(match.group(1))
                        cost = position.Cost(cost_number, self.currency, None, None)
                        price = amount.Amount(rate, self.currency)
                        account_gains = self.account_gains.format(instrument)
                        txn = data.Transaction(
                            meta,
                            date,
                            self.FLAG,
                            None,
                            desc,
                            data.EMPTY_SET,
                            {link},
                            [
                                data.Posting(
                                    self.account_cash, units, None, None, None, None
                                ),
                                data.Posting(
                                    self.account_fees, fees, None, None, None, None
                                ),
                                data.Posting(
                                    account_inst, units_inst, cost, price, None, None
                                ),
                                data.Posting(account_gains, None, None, None, None, None),
                            ],
                        )

                else:
                    logging.error("Unknown row type: %s; skipping", rtype)
                    continue

                entries.append(txn)

        # Insert a final balance check.
        if index:
            entries.append(
                data.Balance(
                    meta,
                    date + datetime.timedelta(days=1),
                    self.account_cash,
                    amount.Amount(D(row["BALANCE"]), self.currency),
                    None,
                    None,
                )
            )

        return entries
