"""UK tax regime — 2026/27 PAYE/NI/CGT/dividend rules.

Rates are baked-in for 2026/27 and held in module-level constants so
they can be patched per-test or upgraded for a future tax year. Only
the *income* side is modelled at year resolution — pension wrapper
contributions and accumulation-phase tax-relief are handled by the
simulator's ISA/SIPP bucket plumbing, not here.

Sources:
- HMRC rates and thresholds 2026-27 (gov.uk/income-tax-rates).
- TCGA 1992 s.10A — 5-year temporary non-residence claw-back.
"""
from __future__ import annotations

from decimal import Decimal

from fire_planner.tax.base import TaxBreakdown, TaxInputs, TaxRegime, apply_brackets

INF = Decimal("Infinity")

# 2026/27 thresholds — frozen by Treasury until at least 2028-04 per the
# autumn 2024 budget. PA tapers above £100k at £1 per £2 → 0 at £125,140.
PERSONAL_ALLOWANCE = Decimal("12570")
PA_TAPER_FLOOR = Decimal("100000")
PA_TAPER_CEILING = Decimal("125140")
BASIC_RATE_BAND = Decimal("37700")
ADDITIONAL_RATE_THRESHOLD = Decimal("125140")

# PAYE income-tax brackets, applied to TAXABLE income (after PA).
# HMRC defines the additional-rate threshold at £125,140 of *taxable*
# income — independent of the PA-taper outcome — so the higher-rate
# band width depends on PA. With full PA the 40% band runs from
# £37,701 to £125,140 (width £87,440); with PA fully tapered the 40%
# band still runs to £125,140 of taxable, just from a lower starting
# gross.
INCOME_TAX_BRACKETS: list[tuple[Decimal, Decimal]] = [
    (BASIC_RATE_BAND, Decimal("0.20")),
    (ADDITIONAL_RATE_THRESHOLD, Decimal("0.40")),
    (INF, Decimal("0.45")),
]

# NI Class 1 employee 2026/27 — annualised. Real-world NI is calculated
# per-period but for retirement modelling annual approximation is fine.
NI_PRIMARY_THRESHOLD = Decimal("12570")
NI_UPPER_EARNINGS_LIMIT = Decimal("50270")
NI_BRACKETS: list[tuple[Decimal, Decimal]] = [
    (NI_UPPER_EARNINGS_LIMIT - NI_PRIMARY_THRESHOLD, Decimal("0.08")),
    (INF, Decimal("0.02")),
]

# Capital gains — Autumn Budget 2024 equalised property + non-property
# rates within each band. £3,000 annual exempt amount.
CGT_ANNUAL_EXEMPTION = Decimal("3000")
CGT_BASIC_RATE = Decimal("0.18")
CGT_HIGHER_RATE = Decimal("0.24")

# Dividend tax 2026/27 — £500 allowance, then 8.75 / 33.75 / 39.35.
DIVIDEND_ALLOWANCE = Decimal("500")
DIVIDEND_BASIC = Decimal("0.0875")
DIVIDEND_HIGHER = Decimal("0.3375")
DIVIDEND_ADDITIONAL = Decimal("0.3935")

# Personal Savings Allowance — only basic and higher rate get any.
PSA_BASIC = Decimal("1000")
PSA_HIGHER = Decimal("500")

# Pension withdrawal — 25% tax-free up to the lump-sum allowance, rest
# taxed as ordinary income. Cap is per-lifetime; we apply it per-year
# because the simulator doesn't track cumulative PCLS yet (all
# withdrawals stay below the cap on the £100k spend assumption).
PCLS_FRACTION = Decimal("0.25")


def taper_personal_allowance(adjusted_net_income: Decimal) -> Decimal:
    """Apply the PA taper: 0 above £125,140, full PA below £100k."""
    if adjusted_net_income <= PA_TAPER_FLOOR:
        return PERSONAL_ALLOWANCE
    if adjusted_net_income >= PA_TAPER_CEILING:
        return Decimal("0")
    reduction = (adjusted_net_income - PA_TAPER_FLOOR) / Decimal("2")
    return max(Decimal("0"), PERSONAL_ALLOWANCE - reduction)


def _income_tax(taxable_income: Decimal) -> Decimal:
    return apply_brackets(taxable_income, INCOME_TAX_BRACKETS)


def _ni(earned_income: Decimal) -> Decimal:
    above_threshold = max(Decimal("0"), earned_income - NI_PRIMARY_THRESHOLD)
    return apply_brackets(above_threshold, NI_BRACKETS)


def _cgt(gains: Decimal, taxable_non_gains_income: Decimal) -> Decimal:
    """Apply CGT — fills the unused basic rate band first, then 24%."""
    after_exemption = max(Decimal("0"), gains - CGT_ANNUAL_EXEMPTION)
    if after_exemption == 0:
        return Decimal("0")
    basic_band_remaining = max(Decimal("0"), BASIC_RATE_BAND - taxable_non_gains_income)
    in_basic = min(after_exemption, basic_band_remaining)
    in_higher = after_exemption - in_basic
    return in_basic * CGT_BASIC_RATE + in_higher * CGT_HIGHER_RATE


def _dividend_tax(dividends: Decimal, taxable_non_div_income: Decimal) -> Decimal:
    """Dividends are stacked on top of other income, so the band
    boundaries depend on what's already used. £500 allowance off the top.
    """
    after_allowance = max(Decimal("0"), dividends - DIVIDEND_ALLOWANCE)
    if after_allowance == 0:
        return Decimal("0")
    basic_band_remaining = max(Decimal("0"), BASIC_RATE_BAND - taxable_non_div_income)
    higher_band_remaining = max(
        Decimal("0"), ADDITIONAL_RATE_THRESHOLD - max(taxable_non_div_income, BASIC_RATE_BAND))

    in_basic = min(after_allowance, basic_band_remaining)
    rest = after_allowance - in_basic
    in_higher = min(rest, higher_band_remaining)
    in_additional = rest - in_higher
    return (in_basic * DIVIDEND_BASIC + in_higher * DIVIDEND_HIGHER +
            in_additional * DIVIDEND_ADDITIONAL)


def _psa_for_band(taxable_non_savings_income: Decimal) -> Decimal:
    """Personal Savings Allowance scales by tax band:
    - basic rate (income within £37,700): £1,000
    - higher rate: £500
    - additional rate (above £125,140 net of PA): £0
    """
    if taxable_non_savings_income <= BASIC_RATE_BAND:
        return PSA_BASIC
    if taxable_non_savings_income <= ADDITIONAL_RATE_THRESHOLD:
        return PSA_HIGHER
    return Decimal("0")


class UkTaxRegime(TaxRegime):
    """UK 2026/27. ISA withdrawals are pre-filtered out (always 0%);
    pension withdrawals get 25% tax-free, the rest is added to earned
    income for PAYE.

    The 5-year Temporary Non-Residence claw-back is the simulator's
    job: when a path returns to the UK within 5y of departure, it
    sums the non-UK regime's pre-tax flows for those years and runs
    them through this regime to compute the recapture. This class
    just computes "tax in a single UK year".
    """
    name = "uk"

    def compute_tax(self, inputs: TaxInputs) -> TaxBreakdown:
        # 25% PCLS, rest taxed as income.
        pcls_tax_free = inputs.pension_withdrawal * PCLS_FRACTION
        pension_taxable = inputs.pension_withdrawal - pcls_tax_free

        ordinary_income = inputs.earned_income + pension_taxable
        adjusted_net_income = ordinary_income + inputs.dividends + inputs.interest
        pa = taper_personal_allowance(adjusted_net_income)
        taxable_ordinary = max(Decimal("0"), ordinary_income - pa)

        income_tax = _income_tax(taxable_ordinary)
        ni = _ni(inputs.earned_income)

        psa = _psa_for_band(taxable_ordinary)
        taxable_interest = max(Decimal("0"), inputs.interest - psa)
        income_tax += apply_brackets(taxable_interest, INCOME_TAX_BRACKETS)

        cgt = _cgt(inputs.capital_gains, taxable_ordinary)
        div_tax = _dividend_tax(inputs.dividends, taxable_ordinary)

        return TaxBreakdown(
            income_tax=income_tax,
            national_insurance=ni,
            capital_gains_tax=cgt,
            dividend_tax=div_tax,
            notes=("uk-2026-27", f"pcls_tax_free={pcls_tax_free}", f"pa_used={pa}"),
        )