from odoo import api, fields, models
from odoo.tools import float_compare, float_is_zero
from collections import defaultdict


class AccountReconcileModel(models.Model):
    _inherit = 'account.reconcile.model'

    set_partner_id = fields.Many2one('res.partner', string='Set Partner')

    """
    Big override/patch to extend the behavior during proposal of 
    write off if the rule has a "set_partner_id" and the statement
    line does not have a partner set.
    """
    def _apply_rules(self, st_lines, excluded_ids=None, partner_map=None):
        ''' Apply criteria to get candidates for all reconciliation models.
        :param st_lines:        Account.bank.statement.lines recordset.
        :param excluded_ids:    Account.move.lines to exclude.
        :param partner_map:     Dict mapping each line with new partner eventually.
        :return:                A dict mapping each statement line id with:
            * aml_ids:      A list of account.move.line ids.
            * model:        An account.reconcile.model record (optional).
            * status:       'reconciled' if the lines has been already reconciled, 'write_off' if the write-off must be
                            applied on the statement line.
        '''
        # This functions uses SQL to compute its results. We need to flush before doing anything more.
        for model_name in ('account.bank.statement', 'account.bank.statement.line', 'account.move', 'account.move.line', 'res.company', 'account.journal', 'account.account'):
            self.env[model_name].flush(self.env[model_name]._fields)

        results = {line.id: {'aml_ids': []} for line in st_lines}

        available_models = self.filtered(lambda m: m.rule_type != 'writeoff_button').sorted()
        aml_ids_to_exclude = set() # Keep track of already processed amls.
        reconciled_amls_ids = set() # Keep track of already reconciled amls.

        # First associate with each rec models all the statement lines for which it is applicable
        lines_with_partner_per_model = defaultdict(lambda: [])
        for st_line in st_lines:

            # Statement lines created in old versions could have a residual amount of zero. In that case, don't try to
            # match anything.
            if not st_line.amount_residual:
                continue

            mapped_partner = (partner_map and partner_map.get(st_line.id) and self.env['res.partner'].browse(partner_map[st_line.id])) or st_line.partner_id

            for rec_model in available_models:
                partner = mapped_partner or rec_model._get_partner_from_mapping(st_line)

                if rec_model._is_applicable_for(st_line, partner):
                    # Customization
                    if not partner and rec_model.set_partner_id and not rec_model.match_partner:
                        partner = rec_model.set_partner_id
                        st_line.partner_id = partner
                    # End Customization

                    lines_with_partner_per_model[rec_model].append((st_line, partner))

        # Execute only one SQL query for each model (for performance)
        matched_lines = self.env['account.bank.statement.line']
        for rec_model in available_models:

            # We filter the lines for this model, in case a previous one has already found something for them
            filtered_st_lines_with_partner = [x for x in lines_with_partner_per_model[rec_model] if x[0] not in matched_lines]

            if not filtered_st_lines_with_partner:
                # No unreconciled statement line for this model
                continue

            all_model_candidates = rec_model._get_candidates(filtered_st_lines_with_partner, excluded_ids)

            for st_line, partner in filtered_st_lines_with_partner:
                candidates = all_model_candidates[st_line.id]
                if candidates:
                    model_rslt, new_reconciled_aml_ids, new_treated_aml_ids = rec_model._get_rule_result(st_line, candidates, aml_ids_to_exclude, reconciled_amls_ids, partner)

                    if model_rslt:
                        # We inject the selected partner (possibly coming from the rec model)
                        model_rslt['partner']= partner

                        results[st_line.id] = model_rslt
                        reconciled_amls_ids |= new_reconciled_aml_ids
                        aml_ids_to_exclude |= new_treated_aml_ids
                        matched_lines += st_line

        return results