"""Post tax forms processor."""

import logging
from typing import cast

from config import app_logger as logger
from src.connectors import ows_abacus_account, ows_abacus_state, ows_payee
from src.constants import AbacusActions, ActionStatuses, DEFAULT_BATCH_SIZE
from src.models import PostTaxFormsInput, UpdateAccountTaxInfo
from src.processors.base import Processor
from src.processors.exceptions import ProcessingError
from src.utils import chunks


class PostTaxFormsProcessor(Processor):
    """
    Processor to post payee`s tax forms data.
    Expects a CSV file with several of the following columns depending on the current tax form type:
        - vendor_id (str)
        - tax_residence_country (str)
        - tax_name (str)
        - tax_classification (str)
        - tin_type (str)
        - tin (str)
        - tax_id_country (str)
        - tax_form_type (str)
        - tax_classification (str)
        - tax_treaty_claim (bool)
        - signed_date (date)
        - type_of_entity (str)
        - lob (str)
        - override (bool)
    The first row is a header.
    """

    _data: dict[int, PostTaxFormsInput]
    _account_id_account_payee_id: dict[int, int]
    _account_id_account_tax_info_id: dict[int, int]
    _account_id_account_payee_tax_form_info_id: dict[int, int]

    def process(self) -> None:
        """Run main logic."""
        self._load_data()
        self._load_payee_ids()
        self._load_account_tax_info()
        self._load_tax_forms()
        self._validate_data()
        self._post_tax_forms()

    def _load_data(self) -> None:
        """Load data from file."""
        try:
            raw_data = list(self.csv_dict_reader)
        except Exception as e:
            raise ProcessingError(f'Unable to read the file: {e}')

        data = {}
        for record in raw_data:
            try:
                post_input = PostTaxFormsInput.model_validate(record)
                data[post_input.account_id] = post_input
            except Exception as e:
                self._add_log(str(e))
        self._data = data

        if not data:
            raise ProcessingError('No items in file to process')

        logger.info(f'Loaded {len(data)} items.')

    def _load_payee_ids(self) -> None:
        """Load account payee IDs by account IDs."""
        account_id_account_payee_id = {}
        for account_ids in chunks(list(self._data.keys()), DEFAULT_BATCH_SIZE):
            account_id_account_payee_id.update(
                ows_abacus_account.get_payees_by_accounts(account_ids)
            )
        self._account_id_account_payee_id = account_id_account_payee_id

    def _load_account_tax_info(self) -> None:
        """Load account tax info by account IDs."""
        account_id_account_tax_info_id = {}
        for account_ids in chunks(list(self._data.keys()), DEFAULT_BATCH_SIZE):
            response = ows_abacus_account.get_account_tax_info_bulk(
                account_ids, limit=DEFAULT_BATCH_SIZE
            )
            for tax_info in response.items:
                account_id_account_tax_info_id[tax_info.account_id] = (
                    tax_info.account_tax_info_id
                )
        self._account_id_account_tax_info_id = account_id_account_tax_info_id

    def _load_tax_forms(self) -> None:
        """Load tax forms via the ows-payee."""
        result = {}
        for chunk in chunks(
            list(self._account_id_account_payee_id.items()), DEFAULT_BATCH_SIZE
        ):
            account_payee_id_account_id = {
                payee_id: account_id for (account_id, payee_id) in chunk
            }
            tax_forms = ows_payee.get_tax_form_info_bulk(
                list(account_payee_id_account_id.keys()), limit=DEFAULT_BATCH_SIZE
            )
            for tax_form_info in tax_forms.items:
                account_id = account_payee_id_account_id[tax_form_info.account_payee_id]
                result[account_id] = tax_form_info.account_payee_tax_form_info_id
        self._account_id_account_payee_tax_form_info_id = result

    def _validate_data(self) -> None:
        validated_data = {}

        for account_id, post_input in self._data.items():
            account_payee_id = self._account_id_account_payee_id.get(account_id)
            if not account_payee_id:
                self._add_log(f'Missing account_payee for account {account_id}')
                continue

            account_tax_info_id = self._account_id_account_tax_info_id.get(account_id)
            if not account_tax_info_id:
                self._add_log(f'Missing account_tax_info for account {account_id}')
                continue

            account_payee_tax_form_info_id = (
                self._account_id_account_payee_tax_form_info_id.get(account_id)
            )
            if account_payee_tax_form_info_id:
                if post_input.override:
                    self._add_log(
                        f'Tax forms exist for account {account_id}. Override',
                        logging.INFO,
                    )
                else:
                    self._add_log(
                        f'Tax forms exist for account {account_id}. Skip',
                        logging.WARNING,
                    )
                    continue

            validated_data[account_id] = post_input.model_copy(
                update={
                    'account_payee_id': account_payee_id,
                    'account_tax_info_id': account_tax_info_id,
                    'account_payee_tax_form_info_id': account_payee_tax_form_info_id,
                }
            )

        self._data = validated_data

    def _post_tax_forms(self) -> None:
        """Post tax form info via the ows-payee."""
        for account_id, post_input in self._data.items():
            try:
                if post_input.override and post_input.account_payee_tax_form_info_id:
                    ows_payee.delete_tax_form_info(
                        post_input.account_payee_tax_form_info_id
                    )
                ows_payee.save_tax_form_info(
                    cast(int, post_input.account_payee_id), post_input.tax_form
                )
                ows_abacus_state.set_state(
                    cast(int, post_input.account_payee_id),
                    AbacusActions.tax_eligibility,
                    ActionStatuses.complete,
                    'Upload tax form from S3 file',
                )
                ows_abacus_account.update_account_tax_info(
                    cast(int, post_input.account_tax_info_id),
                    UpdateAccountTaxInfo(
                        country_of_tax_residence=post_input.country_of_tax_residence,
                        is_tax_treaty_claimed=getattr(
                            post_input.tax_form, 'tax_treaty_claim', None
                        ),
                    ),
                )
            except Exception as e:
                self._add_log(f'{account_id}/{post_input.account_payee_id} {str(e)}')