from odoo import api, fields, models


class ContractContract(models.Model):
    _inherit = "contract.contract"

    payment_mode_id = fields.Many2one(
        comodel_name="account.payment.mode",
        string="Payment Mode",
        domain=[("payment_type", "=", "inbound")],
        index=True,
    )

    @api.onchange("partner_id")
    def on_change_partner_id(self):
        self.payment_mode_id = self.partner_id.customer_payment_mode_id.id

    @api.multi
    def _prepare_invoice(self, date_invoice, journal=None):
        invoice_vals = super(ContractContract, self)._prepare_invoice(
            date_invoice=date_invoice, journal=journal
        )
        if self.payment_mode_id:
            invoice_vals["payment_mode_id"] = self.payment_mode_id.id
            invoice = self.env["account.invoice"].new(invoice_vals)
            invoice._onchange_payment_mode_id()
            invoice_vals = invoice._convert_to_write(invoice._cache)
        return invoice_vals

    @api.model
    def _finalize_invoice_creation(self, invoices):
        """
        This override preserves the payment mode when calling the partner's
        onchange.
        """
        payment_modes_by_invoice = {}
        for invoice in invoices:
            payment_modes_by_invoice[invoice] = invoice.payment_mode_id
        res = super(ContractContract, self)._finalize_invoice_creation(invoices)
        for invoice in invoices:
            invoice.payment_mode_id = payment_modes_by_invoice.get(invoice)
        return res